diff --git a/runner/sandbox.go b/runner/sandbox.go index b9b374c..f8f69e2 100644 --- a/runner/sandbox.go +++ b/runner/sandbox.go @@ -145,7 +145,6 @@ func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx *Context) (* // Execute with 2 args, 1 result if err := state.Call(2, 1); err != nil { - ReleaseActiveConnections(state) return nil, fmt.Errorf("script execution failed: %w", err) } @@ -159,8 +158,6 @@ func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx *Context) (* extractHTTPResponseData(state, response) - ReleaseActiveConnections(state) - return response, nil } diff --git a/runner/sqlite.go b/runner/sqlite.go index 5c41165..2bcd95d 100644 --- a/runner/sqlite.go +++ b/runner/sqlite.go @@ -2,7 +2,6 @@ package runner import ( "context" - "errors" "fmt" "path/filepath" "strings" @@ -16,174 +15,69 @@ import ( luajit "git.sharkk.net/Sky/LuaJIT-to-Go" ) -// SQLiteConnection tracks an active connection -type SQLiteConnection struct { - Conn *sqlite.Conn - Pool *sqlitex.Pool -} +// DbPools maintains database connection pools +var ( + dbPools = make(map[string]*sqlitex.Pool) + poolsMu sync.RWMutex + dataDir string +) -// SQLiteManager handles database connections -type SQLiteManager struct { - mu sync.RWMutex - pools map[string]*sqlitex.Pool - activeConns map[string]*SQLiteConnection - dataDir string -} - -var sqliteManager *SQLiteManager - -// InitSQLite initializes the SQLite manager -func InitSQLite(dataDir string) { - sqliteManager = &SQLiteManager{ - pools: make(map[string]*sqlitex.Pool), - activeConns: make(map[string]*SQLiteConnection), - dataDir: dataDir, - } - logger.Server("SQLite initialized with data directory: %s", dataDir) +// InitSQLite initializes the SQLite subsystem +func InitSQLite(dir string) { + dataDir = dir + logger.Server("SQLite initialized with data directory: %s", dir) } // CleanupSQLite closes all database connections func CleanupSQLite() { - if sqliteManager == nil { - return - } + poolsMu.Lock() + defer poolsMu.Unlock() - sqliteManager.mu.Lock() - defer sqliteManager.mu.Unlock() - - for id, conn := range sqliteManager.activeConns { - if conn.Pool != nil { - conn.Pool.Put(conn.Conn) - } - delete(sqliteManager.activeConns, id) - } - - for name, pool := range sqliteManager.pools { + for name, pool := range dbPools { if err := pool.Close(); err != nil { logger.Error("Failed to close database %s: %v", name, err) } } - sqliteManager.pools = nil - sqliteManager.activeConns = nil + dbPools = make(map[string]*sqlitex.Pool) logger.Debug("SQLite connections closed") } -// ReleaseActiveConnections returns all active connections to their pools -func ReleaseActiveConnections(state *luajit.State) { - if sqliteManager == nil { - return - } - - // Get active connections table from Lua - state.GetGlobal("__active_sqlite_connections") - if !state.IsTable(-1) { - state.Pop(1) - return - } - - sqliteManager.mu.Lock() - defer sqliteManager.mu.Unlock() - - // Iterate through active connections - state.PushNil() // Start iteration - for state.Next(-2) { - if state.IsTable(-1) { - state.GetField(-1, "id") - if state.IsString(-1) { - connID := state.ToString(-1) - if conn, exists := sqliteManager.activeConns[connID]; exists { - if conn.Pool != nil { - conn.Pool.Put(conn.Conn) - } - delete(sqliteManager.activeConns, connID) - } - } - state.Pop(1) // Pop connection id - } - state.Pop(1) // Pop value, leave key for next iteration - } - - // Clear the active connections table - state.PushNil() - state.SetGlobal("__active_sqlite_connections") -} - -// getConnection returns a connection for the database -func getConnection(dbName, connID string) (*sqlite.Conn, error) { - if sqliteManager == nil { - return nil, errors.New("SQLite not initialized") - } - +// getPool returns a connection pool for the database +func getPool(dbName string) (*sqlitex.Pool, error) { // Validate database name dbName = filepath.Base(dbName) if dbName == "" || dbName[0] == '.' { - return nil, errors.New("invalid database name") + return nil, fmt.Errorf("invalid database name") } - // Check for existing connection - sqliteManager.mu.RLock() - conn, exists := sqliteManager.activeConns[connID] + // Check for existing pool + poolsMu.RLock() + pool, exists := dbPools[dbName] if exists { - sqliteManager.mu.RUnlock() - return conn.Conn, nil + poolsMu.RUnlock() + return pool, nil } - sqliteManager.mu.RUnlock() + poolsMu.RUnlock() - // Get or create pool under write lock - sqliteManager.mu.Lock() - defer sqliteManager.mu.Unlock() + // Create new pool under write lock + poolsMu.Lock() + defer poolsMu.Unlock() - // Double-check if a connection was created while waiting for lock - if conn, exists = sqliteManager.activeConns[connID]; exists { - return conn.Conn, nil + // Double-check if a pool was created while waiting for lock + if pool, exists = dbPools[dbName]; exists { + return pool, nil } - // Get or create pool - pool, exists := sqliteManager.pools[dbName] - if !exists { - dbPath := filepath.Join(sqliteManager.dataDir, dbName+".db") - var err error - pool, err = sqlitex.NewPool(dbPath, sqlitex.PoolOptions{}) - if err != nil { - return nil, fmt.Errorf("failed to open database: %w", err) - } - sqliteManager.pools[dbName] = pool - } - - // Get a connection - dbConn, err := pool.Take(context.Background()) + // Create new pool + dbPath := filepath.Join(dataDir, dbName+".db") + pool, err := sqlitex.NewPool(dbPath, sqlitex.PoolOptions{}) if err != nil { - return nil, fmt.Errorf("failed to get connection from pool: %w", err) + return nil, fmt.Errorf("failed to open database: %w", err) } - // Store connection - sqliteManager.activeConns[connID] = &SQLiteConnection{ - Conn: dbConn, - Pool: pool, - } - - return dbConn, nil -} - -// releaseConnection returns a connection to its pool -func releaseConnection(connID string) { - if sqliteManager == nil { - return - } - - sqliteManager.mu.Lock() - defer sqliteManager.mu.Unlock() - - conn, exists := sqliteManager.activeConns[connID] - if !exists { - return - } - - if conn.Pool != nil { - conn.Pool.Put(conn.Conn) - } - delete(sqliteManager.activeConns, connID) + dbPools[dbName] = pool + return pool, nil } // sqlQuery executes a SQL query and returns results @@ -196,15 +90,21 @@ func sqlQuery(state *luajit.State) int { dbName := state.ToString(1) query := state.ToString(2) - connID := fmt.Sprintf("temp_%p", &query) - // Get connection - conn, err := getConnection(dbName, connID) + // Get connection pool + pool, err := getPool(dbName) if err != nil { state.PushString(fmt.Sprintf("sqlite.query: %s", err.Error())) return -1 } - defer releaseConnection(connID) + + // Get a connection from the pool + conn, err := pool.Take(context.Background()) + if err != nil { + state.PushString(fmt.Sprintf("sqlite.query: failed to get connection: %s", err.Error())) + return -1 + } + defer pool.Put(conn) // Create execution options var execOpts sqlitex.ExecOptions @@ -326,15 +226,21 @@ func sqlExec(state *luajit.State) int { dbName := state.ToString(1) query := state.ToString(2) - connID := fmt.Sprintf("temp_%p", &query) - // Get connection - conn, err := getConnection(dbName, connID) + // Get connection pool + pool, err := getPool(dbName) if err != nil { state.PushString(fmt.Sprintf("sqlite.exec: %s", err.Error())) return -1 } - defer releaseConnection(connID) + + // Get a connection from the pool + conn, err := pool.Take(context.Background()) + if err != nil { + state.PushString(fmt.Sprintf("sqlite.exec: failed to get connection: %s", err.Error())) + return -1 + } + defer pool.Put(conn) // Check if parameters are provided hasParams := state.GetTop() >= 3 && !state.IsNil(3)