From 757672c4a3ce8417bbdbe036df9271b7f3185a81 Mon Sep 17 00:00:00 2001 From: "turbolytics.io" Date: Sun, 13 Sep 2026 17:28:30 -0400 Subject: [PATCH] Every statement on the pipeline's DuckDB connection holds the shared lock A window manager failed polls on a healthy pipeline with "Attempting to execute an unsuccessful or closed pending query result". Once on the Bluesky demo at batch 500, and seven times in ten minutes at batch 1 on bluesky.kafka.windowed.yml with the v1.2.0 image. Since #270 a failed poll stops the process, so on main this is a crash at a rate set by batch size. DuckDB closes a pending result the moment another statement runs on its connection. The pipeline, the managers, the progress store and the debug API share one connection and one mutex, and three parties ran statements without the mutex: - The handler reset. Init drops or truncates the batch table after every batch and once at loop start. Unlocked since the Go rewrite (#123), so the race shipped in v1.0.0. At batch 1 it runs once per message. - The progress write. recordProgress released the lock, then ran its UPDATE. Added in #259. - The sqlcommand sink. Its flush runs a drop, an ingest and the user's SQL, as the pipeline sink, a manager sink or the DLQ sink. Each now takes the lock. The sink gets it through sinks.WithConnLock, and sinks.New refuses to build a sqlcommand sink without it, so a new call site that forgets fails at startup instead of racing. No caller flushes a sink while holding the lock, so the sink taking it cannot deadlock. A race has no deterministic reproduction, so the tests check the invariant instead. duckdb.LockChecked wraps a connection and records any statement, commit or rollback that finds the lock free, with the frames that led there. Two tests run single-threaded under it, where a free lock can only mean the caller let go: - core: the consume loop with a handler that touches the connection and a real progress store. Before the fix it named the handler reset at both call sites and the progress write. - run: the real InferredMemBatch handler, sqlcommand sinks in all three roles, the progress store and a tumbling manager, built by the functions the run command calls. Before the fix it named 15 sink statements. Evidence: - Both tests fail on the old code and pass on the new. - Fixing the progress write alone still failed a poll after 2,900 messages on the batch-1 config. That run is how the handler reset was found. - With all three fixed, the same config ran 15,000 messages with no poll failure. bluesky.postgres.windowed.yml ran 12,000 messages with no poll failure and published closed minutes to Postgres through the locked manager sink. - go test -short -race ./... passes. What breaks if this is wrong: a statement still runs unlocked and a manager poll fails, which on main stops the pipeline. A sink that took the lock while its caller held it would deadlock the pipeline instead. Closes #280. --- internal/cli/examples_test.go | 8 +- internal/cli/run/managers.go | 8 +- internal/cli/run/root.go | 28 ++++- internal/cli/run/row_roles_test.go | 3 +- internal/cli/run/shared_conn_test.go | 168 +++++++++++++++++++++++++++ internal/core/lock_test.go | 96 +++++++++++++++ internal/core/turbine.go | 24 +++- internal/duckdb/lockcheck.go | 157 +++++++++++++++++++++++++ internal/sinks/init.go | 21 ++++ internal/sinks/sqlcommand.go | 17 +++ internal/sinks/sqlcommand_test.go | 18 +++ 11 files changed, 535 insertions(+), 13 deletions(-) create mode 100644 internal/cli/run/shared_conn_test.go create mode 100644 internal/core/lock_test.go create mode 100644 internal/duckdb/lockcheck.go diff --git a/internal/cli/examples_test.go b/internal/cli/examples_test.go index b3ea7ed1..3c502de0 100644 --- a/internal/cli/examples_test.go +++ b/internal/cli/examples_test.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -166,7 +167,10 @@ func TestConfigValidation_ExampleConfigsBuildRealComponents(t *testing.T) { // nothing the commands above may have failed to create. Building // it second meant a handler that skipped for a missing ATTACHed // table took the sink check down with it. - _, err = sinks.New(ctx, conf.Pipeline.Sink, conn) + // The run command hands every sink the connection's lock, and a + // sqlcommand sink refuses to build without one. + lock := &sync.Mutex{} + _, err = sinks.New(ctx, conf.Pipeline.Sink, conn, sinks.WithConnLock(lock)) checkBuildError(t, "sink", err) if conf.Tables != nil { @@ -174,7 +178,7 @@ func TestConfigValidation_ExampleConfigsBuildRealComponents(t *testing.T) { if table.Window == nil { continue } - _, err := sinks.New(ctx, table.Window.Sink, conn) + _, err := sinks.New(ctx, table.Window.Sink, conn, sinks.WithConnLock(lock)) checkBuildError(t, "window sink for "+table.Name, err) } } diff --git a/internal/cli/run/managers.go b/internal/cli/run/managers.go index 378d9957..d1d7e64d 100644 --- a/internal/cli/run/managers.go +++ b/internal/cli/run/managers.go @@ -110,10 +110,16 @@ func buildManagedTables( "table %q: opening the window sink's connection", table.Name) } conns = append(conns, sinkConn) + // The sink's connection is the window's alone, and the manager runs + // the sink from one goroutine, so the lock a sqlcommand sink needs + // has no other party. It is a lock of its own rather than the + // pipeline's: taking the pipeline's here would stall the consume + // loop for the length of every window write. sink, err := sinks.New(ctx, table.Window.Sink, sinkConn, sinks.WithMeterProvider(mp), sinks.WithSinkRole("manager"), - sinks.WithRetryEvents(events)) + sinks.WithRetryEvents(events), + sinks.WithConnLock(&sync.Mutex{})) if err != nil { return nil, closeConns, fmt.Errorf("table %q window sink: %w", table.Name, err) } diff --git a/internal/cli/run/root.go b/internal/cli/run/root.go index d018386b..7381917e 100644 --- a/internal/cli/run/root.go +++ b/internal/cli/run/root.go @@ -37,6 +37,7 @@ func newErrorPolicies( ctx context.Context, conf *config.Conf, conn adbc.Connection, + lock *sync.Mutex, mp metric.MeterProvider, events sinks.RetryEvents, ) (core.PipelineErrorPolicies, error) { @@ -63,7 +64,8 @@ func newErrorPolicies( dlqSink, err := sinks.New(ctx, *onError.DLQ, conn, sinks.WithMeterProvider(mp), sinks.WithSinkRole("dlq"), - sinks.WithRetryEvents(events)) + sinks.WithRetryEvents(events), + sinks.WithConnLock(lock)) if err != nil { return policies, fmt.Errorf("pipeline.on_error dlq: %w", err) } @@ -73,6 +75,23 @@ func newErrorPolicies( return policies, nil } +// newPipelineSink builds the pipeline's own sink. A function of its own so the +// shared-connection test builds it exactly as the run command does. +func newPipelineSink( + ctx context.Context, + conf *config.Conf, + conn adbc.Connection, + lock *sync.Mutex, + mp metric.MeterProvider, + events sinks.RetryEvents, +) (core.Sink, error) { + return sinks.New(ctx, conf.Pipeline.Sink, conn, + sinks.WithMeterProvider(mp), + sinks.WithSinkRole(core.SinkRolePipeline), + sinks.WithRetryEvents(events), + sinks.WithConnLock(lock)) +} + func NewCommand() *cobra.Command { var configPath string var maxMsgs int @@ -375,10 +394,7 @@ func NewCommand() *cobra.Command { // The signal context, so a SIGTERM arriving while the sink dials // its destination stops the start instead of waiting it out. - sink, err := sinks.New(ctx, conf.Pipeline.Sink, conn, - sinks.WithMeterProvider(meterProvider), - sinks.WithSinkRole(core.SinkRolePipeline), - sinks.WithRetryEvents(retryEvents)) + sink, err := newPipelineSink(ctx, conf, conn, lock, meterProvider, retryEvents) if err != nil { return err } @@ -401,7 +417,7 @@ func NewCommand() *cobra.Command { }() } - errorPolicies, err := newErrorPolicies(ctx, conf, conn, meterProvider, retryEvents) + errorPolicies, err := newErrorPolicies(ctx, conf, conn, lock, meterProvider, retryEvents) if err != nil { return err } diff --git a/internal/cli/run/row_roles_test.go b/internal/cli/run/row_roles_test.go index 51c6606b..9e57e191 100644 --- a/internal/cli/run/row_roles_test.go +++ b/internal/cli/run/row_roles_test.go @@ -3,6 +3,7 @@ package run import ( "context" "os" + "sync" "testing" "github.com/apache/arrow-adbc/go/adbc" @@ -112,7 +113,7 @@ func TestDLQRowsCarryTheDLQRole(t *testing.T) { DLQ: &config.Sink{Type: "console"}, } - policies, err := newErrorPolicies(context.Background(), conf, nil, mp, sinks.RetryEvents{}) + policies, err := newErrorPolicies(context.Background(), conf, nil, &sync.Mutex{}, mp, sinks.RetryEvents{}) assert.NoError(t, err) assert.That(t, policies.DLQSink != nil) diff --git a/internal/cli/run/shared_conn_test.go b/internal/cli/run/shared_conn_test.go new file mode 100644 index 00000000..6a0f13fa --- /dev/null +++ b/internal/cli/run/shared_conn_test.go @@ -0,0 +1,168 @@ +package run + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/config" + "github.com/turbolytics/sql-flow/internal/core" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/turbolytics/sql-flow/internal/handlers" + "github.com/turbolytics/sql-flow/internal/sinks" + "github.com/zeebo/assert" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.uber.org/zap" +) + +// sliceSource delivers fixed batches and then closes its stream. +type sliceSource struct { + batches [][]core.Message +} + +func (s *sliceSource) Start() error { return nil } +func (s *sliceSource) Commit() error { return nil } +func (s *sliceSource) Close() error { return nil } + +func (s *sliceSource) Stream() <-chan []core.Message { + ch := make(chan []core.Message, len(s.batches)) + for _, b := range s.batches { + ch <- b + } + close(ch) + return ch +} + +func sharedConnCount(t *testing.T, conn adbc.Connection, table string) int64 { + t.Helper() + stmt, err := conn.NewStatement() + assert.NoError(t, err) + defer stmt.Close() + assert.NoError(t, stmt.SetSqlQuery("SELECT count(*) FROM "+table)) + reader, _, err := stmt.ExecuteQuery(context.Background()) + assert.NoError(t, err) + defer reader.Release() + assert.That(t, reader.Next()) + return reader.Record().Column(0).(*array.Int64).Value(0) +} + +// Every party that runs statements on the pipeline's DuckDB connection holds +// the shared lock while it does: the handler, the pipeline sink, the DLQ sink, +// the progress store, and each table manager with its own sink. +// +// DuckDB closes a pending result the moment another statement runs on its +// connection. A party that skips the lock fails whichever query another party +// has in flight, with "Attempting to execute an unsuccessful or closed pending +// query result". #280: the handler reset and the progress write ran outside +// the lock, the Bluesky demo failed a window poll at batch 500, and at batch 1 +// the failure recurred every few minutes. On main a failed poll stops the +// process. +// +// The race has no deterministic reproduction, so this checks the invariant on +// every statement instead. The components are built by the same functions the +// run command calls, over a connection that records any execution finding the +// lock free. Nothing else touches the connection while the loop and the poll +// run, so a free lock means a party let go of it. +// +// It covers the two places that build a sqlcommand sink on the shared +// connection, the pipeline's and the DLQ's. sinks.New refuses to build one +// without the lock; this test is what shows the sink then uses it. The window +// manager is built beside them and polls, but it runs on connections of its +// own since #281, so it is not a party to the shared one: the lock checker +// never sees it, and that is the point. +func TestPipelineSharedConnection_EveryPartyHoldsTheLock(t *testing.T) { + coverage.Covers(t, "manager.window", "sink.sqlcommand", "error.dlq", "handler.inferred_mem") + ctx := context.Background() + + db, raw := rowsTestDB(t) + rowsTestExec(t, raw, "CREATE TABLE win (bucket TIMESTAMPTZ, n BIGINT)") + rowsTestExec(t, raw, "CREATE TABLE published (bucket TIMESTAMPTZ, n BIGINT)") + rowsTestExec(t, raw, "CREATE TABLE dead (error VARCHAR, message VARCHAR, phase VARCHAR, timestamp VARCHAR)") + + lock := &sync.Mutex{} + conn := duckdb.NewLockChecked(raw, lock) + mp := sdkmetric.NewMeterProvider() + + conf := &config.Conf{ + Tables: &config.Tables{ + SQL: []config.TableSQL{{ + Name: "win", + // One-second buckets with no grace: a bucket closes once a + // later one exists. + Window: &config.Window{ + TimeColumn: "bucket", + SizeSeconds: 1, + LateRows: "drop", + PollIntervalSecs: 3600, + EmitSQL: "SELECT bucket, sum(n)::BIGINT AS n FROM closed GROUP BY ALL", + Sink: config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO published SELECT bucket, n FROM sqlflow_sink_batch", + }}, + }, + }}, + }, + } + conf.Pipeline.Handler = config.Handler{ + Type: "handlers.InferredMemBatch", + SQL: "SELECT to_timestamp(a) AS bucket, count(*)::BIGINT AS n FROM batch GROUP BY ALL", + } + conf.Pipeline.Sink = config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO win SELECT bucket, n FROM sqlflow_sink_batch", + }} + conf.Pipeline.OnError = &config.OnError{ + Policy: "DLQ", + DLQ: &config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO dead SELECT * FROM sqlflow_sink_batch", + }}, + } + + // Startup runs before any second goroutine exists, and the run command + // does it without the lock. Holding it here keeps the check about the + // running pipeline. + progress := core.NewProgressStore(conn) + lock.Lock() + assert.NoError(t, progress.Init(ctx)) + assert.NoError(t, initWindowStores(ctx, conf, conn)) + lock.Unlock() + + sink, err := newPipelineSink(ctx, conf, conn, lock, mp, sinks.RetryEvents{}) + assert.NoError(t, err) + policies, err := newErrorPolicies(ctx, conf, conn, lock, mp, sinks.RetryEvents{}) + assert.NoError(t, err) + handler, err := handlers.New(conn, conf.Pipeline.Handler, zap.NewNop()) + assert.NoError(t, err) + managed, closeConns, err := buildManagedTables(ctx, conf, db, zap.NewNop(), mp, + nil, sinks.RetryEvents{}) + assert.NoError(t, err) + defer closeConns() + assert.Equal(t, 1, len(managed)) + + msg := func(v string) core.Message { return core.Message{Value: []byte(v)} } + src := &sliceSource{batches: [][]core.Message{ + {msg(`{"a": 1}`), msg(`{"a": 1}`)}, + {msg(`not json`), msg(`{"a": 2}`)}, + {msg(`{"a": 2}`), msg(`{"a": 3}`)}, + }} + tb := core.NewTurbine(src, handler, sink, 2, time.Hour, lock, policies, + core.WithProgressStore(progress)) + + _, err = tb.ConsumeLoop(ctx, 6) + assert.NoError(t, err) + assert.NoError(t, managed[0].Poll(ctx)) + + // Every party ran at least once, or a clean result proves nothing. The + // window closed the two buckets a later one exists for, and the newest + // stays open. + assert.Equal(t, int64(2), sharedConnCount(t, raw, "published")) + assert.Equal(t, int64(1), sharedConnCount(t, raw, "dead")) + + if v := conn.Violations(); len(v) > 0 { + t.Fatalf("%d statements ran on the shared connection without the lock:\n%s", + len(v), strings.Join(v, "\n")) + } +} diff --git a/internal/core/lock_test.go b/internal/core/lock_test.go new file mode 100644 index 00000000..8c2cce4a --- /dev/null +++ b/internal/core/lock_test.go @@ -0,0 +1,96 @@ +package core + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/zeebo/assert" +) + +// connHandler is a fakeHandler that touches the connection the way the real +// handlers do: Init resets a batch table and Invoke queries one. The pipeline +// decides whether either runs under the lock, and the fake handler that +// touches nothing hides that decision. +type connHandler struct { + fakeHandler + conn adbc.Connection +} + +func (h *connHandler) exec(ctx context.Context, sql string) error { + stmt, err := h.conn.NewStatement() + if err != nil { + return err + } + defer stmt.Close() + if err := stmt.SetSqlQuery(sql); err != nil { + return err + } + _, err = stmt.ExecuteUpdate(ctx) + return err +} + +func (h *connHandler) Init(ctx context.Context) error { + if err := h.fakeHandler.Init(ctx); err != nil { + return err + } + return h.exec(ctx, "DROP TABLE IF EXISTS batch") +} + +func (h *connHandler) Invoke(ctx context.Context) (arrow.Table, error) { + if err := h.exec(ctx, "CREATE OR REPLACE TABLE batch AS SELECT 1 AS n"); err != nil { + return nil, err + } + return h.fakeHandler.Invoke(ctx) +} + +// Every statement the pipeline runs on the shared connection runs under the +// shared lock. The table managers collect on the same connection under that +// lock, and DuckDB closes a pending result the moment another statement runs +// on its connection, so a statement outside the lock fails whichever poll is +// in flight. #280: the handler reset after every batch and the progress write +// both ran outside the lock, and the Bluesky demo logged "closed pending +// query result" once at batch 500 and seven times in ten minutes at batch 1. +// +// The race cannot be reproduced on demand. The invariant can be checked on +// every statement: the wrapper records any execution that finds the lock +// free. The loop runs with no other goroutine touching the connection, so a +// free lock means the pipeline itself let go of it. +func TestCoreConsumeLoop_EveryStatementOnTheSharedConnectionHoldsTheLock(t *testing.T) { + coverage.Covers(t, "state.durability") + ctx := context.Background() + db, err := duckdb.OpenPath(ctx, "") + assert.NoError(t, err) + defer db.Close() + raw, err := db.Connect(ctx) + assert.NoError(t, err) + defer raw.Close() + + lock := &sync.Mutex{} + conn := duckdb.NewLockChecked(raw, lock) + + // Startup runs before the managers exist, and root.go does it without + // the lock. Holding it here keeps the check about the running pipeline. + progress := NewProgressStore(conn) + lock.Lock() + assert.NoError(t, progress.Init(ctx)) + lock.Unlock() + + src := &fakeSource{batches: [][]Message{messages(4), messages(4), messages(4)}} + tb := NewTurbine(src, &connHandler{conn: conn}, &fakeSink{}, 4, time.Second, + lock, PipelineErrorPolicies{}, WithProgressStore(progress)) + stats, err := tb.ConsumeLoop(ctx, 12) + assert.NoError(t, err) + assert.Equal(t, int64(12), stats.MessagesConsumed()) + + if v := conn.Violations(); len(v) > 0 { + t.Fatalf("%d statements ran on the shared connection without the lock:\n%s", + len(v), strings.Join(v, "\n")) + } +} diff --git a/internal/core/turbine.go b/internal/core/turbine.go index a16f22cd..fbf60f63 100644 --- a/internal/core/turbine.go +++ b/internal/core/turbine.go @@ -388,7 +388,13 @@ func (t *Turbine) recordProgress(ctx context.Context) { return } now := time.Now().UTC() + // Held through the write below. The table managers collect on this same + // connection under this lock, and DuckDB closes a pending result the + // moment another statement runs on its connection. A write outside the + // lock failed whichever poll was in flight with "closed pending query + // result", at a rate set by batch size (#280). t.lock.Lock() + defer t.lock.Unlock() p := Progress{LastCommit: now, Messages: t.stats.MessagesConsumed()} if !t.arrivedAt.IsZero() { p.LastArrival = t.arrivedAt @@ -413,7 +419,6 @@ func (t *Turbine) recordProgress(ctx context.Context) { t.progressWrittenAt = now t.writtenArrival = t.arrivedAt } - t.lock.Unlock() // The snapshot above is exact and free. The table is not: one UPDATE // through ADBC measures about 112 microseconds, and a batch of 5000 at a @@ -594,7 +599,8 @@ func (t *Turbine) ConsumeLoop(ctx context.Context, maxMsgs int) (stats *Stats, e t.stats.StartTime = time.Now().UTC() t.stats.SetNumMessagesConsumed(0) - if err := t.handler.Init(ctx); err != nil { + // Under the lock: the managers are already polling this connection (#280). + if err := t.initHandler(ctx); err != nil { return nil, err } @@ -1033,6 +1039,18 @@ func (t *Turbine) commitSource() error { return t.source.Commit() } +// initHandler resets the handler under the lock. Init drops or truncates the +// batch table on the shared connection, and the table managers poll that +// connection from their own goroutines. DuckDB closes a pending result the +// moment another statement runs on its connection, so an unlocked reset +// failed whichever collect was in flight with "closed pending query result", +// once per batch at batch size 1 (#280). +func (t *Turbine) initHandler(ctx context.Context) error { + t.lock.Lock() + defer t.lock.Unlock() + return t.handler.Init(ctx) +} + // rollbackState discards this batch's uncommitted state writes. Used on the // paths that fail before commitState is reached. func (t *Turbine) rollbackState(ctx context.Context) { @@ -1272,7 +1290,7 @@ func (t *Turbine) processBatch(ctx context.Context, numBatchMessages int) error } i0 := time.Now() - err = t.handler.Init(ctx) + err = t.initHandler(ctx) t.recordPhase(ctx, phaseHandlerInit, time.Since(i0)) if err != nil { diff --git a/internal/duckdb/lockcheck.go b/internal/duckdb/lockcheck.go new file mode 100644 index 00000000..7f7c9557 --- /dev/null +++ b/internal/duckdb/lockcheck.go @@ -0,0 +1,157 @@ +package duckdb + +import ( + "context" + "fmt" + "runtime" + "strings" + "sync" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow/array" +) + +// LockChecked wraps a connection so that every statement executed on it, and +// every commit or rollback, must run while the given lock is held. It is a +// test aid for the one rule the pipeline has about its DuckDB connection. +// +// The pipeline, the table managers, the progress store and the debug API all +// share one connection, serialized by one mutex. DuckDB closes a pending +// result on a connection the moment another statement runs on it, so a +// statement that skips the lock does not deadlock or corrupt anything: it +// fails whatever query was in flight, with "Attempting to execute an +// unsuccessful or closed pending query result". The failure lands on the +// other party, at a rate set by traffic, which is why #280 shipped and ran +// for two days before a batch size of one made it reproducible. +// +// A race like that has no deterministic reproduction. The invariant does: +// at the moment a statement executes, the lock is held. TryLock succeeding +// means nobody held it, which is a violation whoever the caller is. TryLock +// failing means someone held it, and this wrapper cannot tell whether that +// someone is the caller; a concurrent holder hides the violation for that +// one call. So run the pipeline single-threaded under the wrapper, and every +// unlocked statement is caught on its first execution. +type LockChecked struct { + adbc.Connection + lock *sync.Mutex + + mu sync.Mutex + violations []string +} + +// NewLockChecked returns conn wrapped so that executing on it without lock +// held is recorded. Wrap after startup: the tables are created before any +// second goroutine exists, and root.go creates them without the lock. +func NewLockChecked(conn adbc.Connection, lock *sync.Mutex) *LockChecked { + return &LockChecked{Connection: conn, lock: lock} +} + +// Violations lists every execution that ran without the lock, each with the +// operation and the frames that led to it. +func (c *LockChecked) Violations() []string { + c.mu.Lock() + defer c.mu.Unlock() + return append([]string(nil), c.violations...) +} + +func (c *LockChecked) check(op string) { + if !c.lock.TryLock() { + return + } + c.lock.Unlock() + + pcs := make([]uintptr, 12) + n := runtime.Callers(3, pcs) + frames := runtime.CallersFrames(pcs[:n]) + var where []string + for { + f, more := frames.Next() + if !strings.HasSuffix(f.File, "lockcheck.go") { + where = append(where, fmt.Sprintf("%s:%d", f.Function, f.Line)) + } + if !more || len(where) == 6 { + break + } + } + + c.mu.Lock() + c.violations = append(c.violations, op+" without the lock\n\t"+strings.Join(where, "\n\t")) + c.mu.Unlock() +} + +func (c *LockChecked) Commit(ctx context.Context) error { + c.check("Commit") + return c.Connection.Commit(ctx) +} + +func (c *LockChecked) Rollback(ctx context.Context) error { + c.check("Rollback") + return c.Connection.Rollback(ctx) +} + +func (c *LockChecked) NewStatement() (adbc.Statement, error) { + stmt, err := c.Connection.NewStatement() + if err != nil { + return nil, err + } + return &lockCheckedStatement{Statement: stmt, conn: c}, nil +} + +// SetOption forwards to the wrapped connection's post-init options, which +// root.go uses to turn autocommit off for a state path. +func (c *LockChecked) SetOption(key, val string) error { + po, ok := c.Connection.(adbc.PostInitOptions) + if !ok { + return fmt.Errorf("lockcheck: wrapped connection does not support SetOption") + } + return po.SetOption(key, val) +} + +type lockCheckedStatement struct { + adbc.Statement + conn *LockChecked + sql string +} + +func (s *lockCheckedStatement) SetSqlQuery(query string) error { + s.sql = query + return s.Statement.SetSqlQuery(query) +} + +func (s *lockCheckedStatement) ExecuteQuery(ctx context.Context) (array.RecordReader, int64, error) { + s.conn.check("ExecuteQuery " + firstLine(s.sql)) + reader, n, err := s.Statement.ExecuteQuery(ctx) + if err != nil { + return nil, n, err + } + return &lockCheckedReader{RecordReader: reader, stmt: s}, n, nil +} + +func (s *lockCheckedStatement) ExecuteUpdate(ctx context.Context) (int64, error) { + s.conn.check("ExecuteUpdate " + firstLine(s.sql)) + return s.Statement.ExecuteUpdate(ctx) +} + +// lockCheckedReader checks each Next. A result is streamed, so a reader +// drained after the lock is released is as exposed as a statement executed +// without it. +type lockCheckedReader struct { + array.RecordReader + stmt *lockCheckedStatement +} + +func (r *lockCheckedReader) Next() bool { + r.stmt.conn.check("Next " + firstLine(r.stmt.sql)) + return r.RecordReader.Next() +} + +func firstLine(sql string) string { + sql = strings.TrimSpace(sql) + if i := strings.IndexByte(sql, '\n'); i >= 0 { + sql = sql[:i] + } + if len(sql) > 60 { + sql = sql[:60] + } + return sql +} diff --git a/internal/sinks/init.go b/internal/sinks/init.go index 5dedef6b..0bf95869 100644 --- a/internal/sinks/init.go +++ b/internal/sinks/init.go @@ -3,6 +3,7 @@ package sinks import ( "context" "sort" + "sync" "github.com/apache/arrow-adbc/go/adbc" "github.com/apache/arrow-go/v18/arrow" @@ -31,6 +32,14 @@ type options struct { meterProvider metric.MeterProvider role string retryEvents RetryEvents + connLock *sync.Mutex +} + +// WithConnLock supplies the lock that serializes the pipeline's DuckDB +// connection. The sqlcommand sink runs its statements on that connection, so +// New refuses to build one without it. See SQLCommandSink.connLock. +func WithConnLock(lock *sync.Mutex) Option { + return func(o *options) { o.connLock = lock } } // RetryEvents is told when a sink's retry ladder runs. Retry fires per failed @@ -94,6 +103,18 @@ func New(ctx context.Context, sink config.Sink, conn adbc.Connection, opts ...Op return nil, err } + // Checked here rather than in the builder, so every caller of New is held + // to it: the pipeline, each table manager and the DLQ build their sinks in + // three different places, and one of them forgetting the lock is a race + // that no single-threaded test of that sink would see. + if s, ok := built.(*SQLCommandSink); ok { + if o.connLock == nil { + return nil, errs.New(errs.CodeSinkInvalid, + "sink: sqlcommand shares the pipeline's DuckDB connection and needs its lock") + } + s.connLock = o.connLock + } + // Checked before the pipeline consumes anything, so a destination that is // not there fails the start rather than the first flush. if err := probe(ctx, built); err != nil { diff --git a/internal/sinks/sqlcommand.go b/internal/sinks/sqlcommand.go index ef97945d..d0ab2878 100644 --- a/internal/sinks/sqlcommand.go +++ b/internal/sinks/sqlcommand.go @@ -26,6 +26,12 @@ type SQLCommandSink struct { sql string substitutions []config.SQLCommandSubstitution + // connLock serializes conn with the pipeline, the table managers and the + // debug API, which all run statements on it. New sets it. Nil leaves + // serialization to the caller, which only a test driving the sink alone + // can promise. + connLock *sync.Mutex + mu sync.Mutex tables []arrow.Table } @@ -93,6 +99,17 @@ func (s *SQLCommandSink) requeue(tables []arrow.Table) { // send is one delivery attempt. It releases nothing: whether these batches can // be dropped is the caller's decision, and it depends on this error. func (s *SQLCommandSink) send(ctx context.Context, tables []arrow.Table) error { + // Held across the drop, the ingest and the user's SQL. The handler and the + // table managers run statements on this connection from other goroutines, + // and DuckDB closes a pending result the moment another statement runs on + // its connection, so an unlocked flush failed whatever query they had in + // flight (#280). Nobody calls Flush holding this lock: the pipeline and the + // managers release it before they flush. + if s.connLock != nil { + s.connLock.Lock() + defer s.connLock.Unlock() + } + if err := s.materialize(ctx, tables); err != nil { return err } diff --git a/internal/sinks/sqlcommand_test.go b/internal/sinks/sqlcommand_test.go index f27a8c0a..bd98593b 100644 --- a/internal/sinks/sqlcommand_test.go +++ b/internal/sinks/sqlcommand_test.go @@ -8,6 +8,7 @@ package sinks import ( "context" "strings" + "sync" "testing" "github.com/apache/arrow-adbc/go/adbc" @@ -109,6 +110,23 @@ func TestSinkSqlcommand_NewRequiresSQL(t *testing.T) { assert.Error(t, err) } +// The sink runs statements on the pipeline's shared connection, so New refuses +// to build one without the lock that serializes it. A call site that forgets +// the option then fails at startup, not with a "closed pending query result" +// in some other party's query at a rate set by traffic (#280). +func TestSinkSqlcommand_NewRefusesWithoutTheConnectionLock(t *testing.T) { + coverage.Covers(t, "sink.sqlcommand") + conn := newSinkTestConn(t) + conf := config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{SQL: "SELECT 1"}} + + _, err := New(context.Background(), conf, conn) + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "needs its lock")) + + _, err = New(context.Background(), conf, conn, WithConnLock(&sync.Mutex{})) + assert.NoError(t, err) +} + // The batch reaches the user's SQL under the name the Python sink registered // it as. A different name and every shipped sqlcommand config stops working. func TestSinkSqlcommand_RunsSQLAgainstTheBatch(t *testing.T) {