remove all connection tracking
This commit is contained in:
parent
8f9a9da5a1
commit
0581d72065
@ -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
|
||||
}
|
||||
|
||||
|
||||
204
runner/sqlite.go
204
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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user