Merge pull request #38 from chdb-io/concurrent-connections

Support multiple concurrent connections per process

Author: wudidapaopao

README.md

@@ -131,6 +131,28 @@ func main() {
 }
 ```
 
+#### Concurrency
+
+chDB runs a single embedded engine per process bound to one data path, but that
+engine accepts multiple connections that execute queries concurrently. The
+`database/sql` driver opens an independent native chDB connection per pooled
+connection, so you can scale read/write parallelism with `SetMaxOpenConns`:
+
+```go
+db, err := sql.Open("chdb", "session=/path/to/data")
+if err != nil {
+        log.Fatal(err)
+}
+defer db.Close()
+
+// Each pooled connection is its own native chDB connection to the same data
+// path, so queries run in parallel instead of serializing on one connection.
+db.SetMaxOpenConns(8)
+```
+
+All connections in a process must share the same data path; opening a second,
+different data path while connections are still open returns an error.
+
 ### Golang API docs
 
 - See [lowApi.md](lowApi.md) for the low level APIs.

chdb-purego/chdb.go

@@ -112,9 +112,14 @@ func newChdbConn(conn *chdb_connection) ChdbConn {
 }
 
 // Close implements ChdbConn.
+//
+// Close is idempotent: it nils the underlying handle after freeing it so a
+// second Close (or a Close racing a Query that checks for a nil handle) does
+// not double-free the native connection.
 func (c *connection) Close() {
 	if c.conn != nil {
 		chdbCloseConn(c.conn)
+		c.conn = nil
 	}
 }
 
@@ -182,9 +187,10 @@ func (c *connection) Ready() bool {
 //   - argc = 2, argv = []string{"--path=/tmp/chdb", "--readonly=1"}
 //
 // Important:
-//   - There can be only one session at a time. If you want to create a new session, you need to close the existing one.
-//   - Creating a new session will close the existing one.
-//   - You need to ensure that the path exists before creating a new session. Or you can use NewConnectionFromConnString.
+//   - chDB supports only one data path per process. Multiple connections to the
+//     same path can be open at once and execute queries concurrently; connecting
+//     to a different path while connections are still open returns an error.
+//   - You need to ensure that the path exists before creating a new connection. Or you can use NewConnectionFromConnString.
 func NewConnection(argc int, argv []string) (ChdbConn, error) {
 	var new_argv []string
 	if (argc > 0 && argv[0] != "clickhouse") || argc == 0 {
@@ -274,8 +280,9 @@ func NewConnection(argc int, argv []string) (ChdbConn, error) {
 //	- "mode=ro" would be "--readonly=1" for clickhouse (read-only mode)
 //
 // Important:
-//   - There can be only one session at a time. If you want to create a new session, you need to close the existing one.
-//   - Creating a new session will close the existing one.
+//   - chDB supports only one data path per process. Multiple connections to the
+//     same path can be open at once and execute queries concurrently; connecting
+//     to a different path while connections are still open returns an error.
 func NewConnectionFromConnString(conn_string string) (ChdbConn, error) {
 	if conn_string == "" || conn_string == ":memory:" {
 		return NewConnection(0, []string{})

chdb-purego/stress_test.go

@@ -184,3 +184,60 @@ func runStress(conn ChdbConn, id, depth int, queries, failures *atomic.Uint64, s
 		queries.Add(1)
 	}
 }
+
+// TestMultiConnectionStress opens several connections to the same (in-memory)
+// data path and drives queries on all of them concurrently. It guards the
+// multi-connection path (refcounted EmbeddedServer + per-connection clients)
+// against crashes and regressions in the issue-#30 signal handling under
+// concurrent connect/query load. Skipped under `go test -short`.
+func TestMultiConnectionStress(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping stress test in -short mode")
+	}
+
+	const (
+		duration = 5 * time.Second
+		nConns   = 4
+		gPerConn = 4
+		recurse  = 16
+	)
+
+	conns := make([]ChdbConn, 0, nConns)
+	for i := 0; i < nConns; i++ {
+		c, err := NewConnectionFromConnString(":memory:")
+		if err != nil {
+			t.Fatalf("connect %d: %v", i, err)
+		}
+		defer c.Close()
+		conns = append(conns, c)
+	}
+
+	var (
+		wg       sync.WaitGroup
+		queries  atomic.Uint64
+		failures atomic.Uint64
+		stop     atomic.Bool
+	)
+
+	for i := 0; i < nConns; i++ {
+		for g := 0; g < gPerConn; g++ {
+			wg.Add(1)
+			go func(c ChdbConn) {
+				defer wg.Done()
+				runStress(c, 0, recurse, &queries, &failures, &stop)
+			}(conns[i])
+		}
+	}
+
+	time.Sleep(duration)
+	stop.Store(true)
+	wg.Wait()
+
+	t.Logf("conns=%d goroutines=%d queries=%d failures=%d (%.0f qps)",
+		nConns, nConns*gPerConn, queries.Load(), failures.Load(),
+		float64(queries.Load())/duration.Seconds())
+
+	if failures.Load() != 0 {
+		t.Fatalf("%d queries failed under multi-connection stress", failures.Load())
+	}
+}

chdb.md

@@ -28,7 +28,7 @@ import "github.com/chdb-io/chdb-go/chdb"
 func Query(queryStr string, outputFormats ...string) (result chdbpurego.ChdbResult, err error)
 ```
 
-Query calls query\_conn with a default in\-memory session and default output format of "CSV" if not provided.
+Query runs a one\-shot query and returns the materialized result \(default output format "CSV"\). chDB allows only one data path per process, so if a session is already open this helper attaches to that session's data path; otherwise it uses an in\-memory database.
 
 <a name="QueryStream"></a>
 ## func [QueryStream](<https://github.com/s0und0fs1lence/chdb-go/blob/main/chdb/wrapper.go#L23>)
@@ -37,7 +37,7 @@ Query calls query\_conn with a default in\-memory session and default output for
 func QueryStream(queryStr string, outputFormats ...string) (result chdbpurego.ChdbStreamResult, err error)
 ```
 
-Query calls query\_conn with a default in\-memory session and default output format of "CSV" if not provided.
+QueryStream is like Query but returns a streaming result that can be read in chunks, for large datasets that should not be fully materialized in memory. Like Query, it attaches to an already\-open session's data path, or uses an in\-memory database when none is open.
 
 <a name="Session"></a>
 ## type [Session](<https://github.com/s0und0fs1lence/chdb-go/blob/main/chdb/session.go#L14-L19>)
@@ -57,7 +57,7 @@ type Session struct {
 func NewSession(paths ...string) (*Session, error)
 ```
 
-NewSession creates a new session with the given path. If path is empty, a temporary directory is created. Note: The temporary directory is removed when Close is called.
+NewSession creates a new session with the given path. If path is empty, the session reuses an already\-open data path, or creates a temporary directory when none is open. Multiple sessions can be open at once as long as they share the same data path \(each owns an independent native connection, so they can run queries in parallel\); opening a session on a different path while another is still open returns an error. The temporary directory is removed when the last session using it is closed.
 
 <a name="Session.Cleanup"></a>
 ### func \(\*Session\) [Cleanup](<https://github.com/s0und0fs1lence/chdb-go/blob/main/chdb/session.go#L86>)

chdb/driver/driver.go

@@ -182,23 +182,43 @@ type connector struct {
 	bufferSize  int
 	isStreaming bool
 	useUnsafe   bool
-	session     *chdb.Session
+	connStr     string
+	keeper      *chdb.Session
 }
 
 // Connect returns a connection to a database.
 func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	if c.driverType == INVALID {
 		return nil, fmt.Errorf("DriverType not supported")
 	}
+	// Each database/sql connection gets its own native chDB connection to the
+	// shared data path, so a pool of connections (MaxOpenConns > 1) yields real
+	// parallel query execution instead of serializing on a single connection.
+	session, err := chdb.NewSession(c.connStr)
+	if err != nil {
+		return nil, err
+	}
 	cc := &conn{
-		udfPath: c.udfPath, session: c.session,
+		udfPath: c.udfPath, session: session,
 		driverType: c.driverType, bufferSize: c.bufferSize,
 		useUnsafe: c.useUnsafe, isStreaming: c.isStreaming,
 	}
 	cc.SetupQueryFun()
 	return cc, nil
 }
 
+// Close releases the connector's keeper session. database/sql calls this from
+// DB.Close() because the connector implements io.Closer. Dropping the keeper
+// reference lets a registry-owned temp directory be removed once all pooled
+// connections are closed as well.
+func (c *connector) Close() error {
+	if c.keeper != nil {
+		c.keeper.Close()
+		c.keeper = nil
+	}
+	return nil
+}
+
 // Driver returns the underying Driver of the connector,
 // compatibility with the Driver method on sql.DB
 func (c *connector) Driver() driver.Driver { return Driver{} }
@@ -221,13 +241,6 @@ func parseConnectStr(str string) (ret map[string]string, err error) {
 }
 func NewConnect(opts map[string]string) (ret *connector, err error) {
 	ret = &connector{}
-	sessionPath, ok := opts[sessionOptionKey]
-	if ok {
-		ret.session, err = chdb.NewSession(sessionPath)
-		if err != nil {
-			return nil, err
-		}
-	}
 	driverType, ok := opts[driverTypeKey]
 	if ok {
 		ret.driverType = parseDriverType(driverType)
@@ -256,13 +269,18 @@ func NewConnect(opts map[string]string) (ret *connector, err error) {
 	if ok {
 		ret.udfPath = udfPath
 	}
-	if ret.session == nil {
 
-		ret.session, err = chdb.NewSession()
-		if err != nil {
-			return nil, err
-		}
+	// Open a "keeper" session that pins the data path (and any temp directory)
+	// for the lifetime of this connector. Each pooled connection then opens its
+	// own session on the same path; the keeper guarantees the engine and temp
+	// dir survive pool churn (when the live connection count briefly hits zero).
+	sessionPath := opts[sessionOptionKey] // "" when not provided
+	ret.keeper, err = chdb.NewSession(sessionPath)
+	if err != nil {
+		return nil, err
 	}
+	ret.connStr = ret.keeper.ConnStr()
+
 	ret.isStreaming = ret.driverType.SupportStreaming()
 	return
 }
@@ -275,7 +293,25 @@ func (d Driver) Open(name string) (driver.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	return cc.Connect(context.Background())
+	c, err := cc.Connect(context.Background())
+	if err != nil {
+		// Connect failed; release the keeper session NewConnect just opened so
+		// it does not leak (database/sql would normally own and close cc).
+		if closer, ok := cc.(*connector); ok {
+			_ = closer.Close()
+		}
+		return nil, err
+	}
+	// On the sql.Open path, database/sql keeps the connector and calls
+	// connector.Close() on db.Close(). A direct Driver.Open caller discards the
+	// connector, so tie the keeper session's lifetime to the returned
+	// connection: closing the conn also releases the keeper.
+	if cn, ok := c.(*conn); ok {
+		if cnr, ok := cc.(*connector); ok {
+			cn.connector = cnr
+		}
+	}
+	return c, nil
 }
 
 // OpenConnector expects the same format as driver.Open
@@ -294,6 +330,10 @@ type conn struct {
 	useUnsafe   bool
 	isStreaming bool
 	session     *chdb.Session
+	// connector is set only on the legacy Driver.Open path (not the sql.Open
+	// path, where database/sql owns and closes the connector). When set, Close
+	// also releases the connector's keeper session.
+	connector *connector
 
 	QueryFun  queryHandle
 	streamFun queryStream
@@ -312,6 +352,16 @@ func prepareValues(values []driver.Value) []driver.NamedValue {
 }
 
 func (c *conn) Close() error {
+	if c.session != nil {
+		c.session.Close()
+		c.session = nil
+	}
+	// Only set on the Driver.Open path; releases the keeper that pins the data
+	// path/temp dir for the connector.
+	if c.connector != nil {
+		_ = c.connector.Close()
+		c.connector = nil
+	}
 	return nil
 }
 

chdb/driver/driver_open_test.go

@@ -0,0 +1,61 @@
+package chdbdriver
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/chdb-io/chdb-go/chdb"
+)
+
+// TestDriverOpenNoKeeperLeak verifies that the legacy Driver.Open path does not
+// leak the connector's keeper session. database/sql owns the connector on the
+// sql.Open path and calls connector.Close(), but a direct Driver.Open(name)
+// discards the connector, so closing the returned conn must release the keeper
+// too. Pre-fix the keeper session (a native connection + a registry refcount)
+// leaks on every Open call.
+func TestDriverOpenNoKeeperLeak(t *testing.T) {
+	baseline := chdb.ActiveSessionRefs()
+
+	c, err := Driver{}.Open("session=" + session.ConnStr())
+	if err != nil {
+		t.Fatalf("Driver.Open failed: %s", err)
+	}
+	if err := c.Close(); err != nil {
+		t.Fatalf("conn.Close failed: %s", err)
+	}
+
+	if got := chdb.ActiveSessionRefs(); got != baseline {
+		t.Fatalf("Driver.Open leaked sessions: refs=%d, baseline=%d (keeper not released on conn.Close)", got, baseline)
+	}
+}
+
+// TestDbCloseReleasesRefs verifies the database/sql path balances refcounts:
+// opening a *sql.DB, running queries across multiple pooled connections, and
+// closing it returns the registry to its baseline (keeper + per-conn sessions
+// all released).
+func TestDbCloseReleasesRefs(t *testing.T) {
+	baseline := chdb.ActiveSessionRefs()
+
+	db, err := sql.Open("chdb", "session="+session.ConnStr())
+	if err != nil {
+		t.Fatalf("open db failed: %s", err)
+	}
+	db.SetMaxOpenConns(4)
+
+	for i := 0; i < 8; i++ {
+		var n int
+		if err := db.QueryRow("SELECT count() FROM numbers(10)").Scan(&n); err != nil {
+			t.Fatalf("query failed: %s", err)
+		}
+		if n != 10 {
+			t.Fatalf("got %d want 10", n)
+		}
+	}
+
+	if err := db.Close(); err != nil {
+		t.Fatalf("db.Close failed: %s", err)
+	}
+	if got := chdb.ActiveSessionRefs(); got != baseline {
+		t.Fatalf("db.Close left refs leaked: refs=%d, baseline=%d", got, baseline)
+	}
+}