diff --git a/server/cmd/gram/deps.go b/server/cmd/gram/deps.go index 7632d32c76a..d22628e3545 100644 --- a/server/cmd/gram/deps.go +++ b/server/cmd/gram/deps.go @@ -60,6 +60,7 @@ import ( bgtriggers "github.com/speakeasy-api/gram/server/internal/background/triggers" "github.com/speakeasy-api/gram/server/internal/billing" "github.com/speakeasy-api/gram/server/internal/cache" + "github.com/speakeasy-api/gram/server/internal/clickhouseclient" "github.com/speakeasy-api/gram/server/internal/conv" "github.com/speakeasy-api/gram/server/internal/email" "github.com/speakeasy-api/gram/server/internal/encryption" @@ -161,6 +162,11 @@ func newGuardianPolicy(c *cli.Context, logger *slog.Logger, tracerProvider trace return policy, nil } +const ( + clickhouseSocketReadTimeout = 70 * time.Second + clickhouseReadOperationTimeout = 75 * time.Second +) + func newClickhouseClient(ctx context.Context, logger *slog.Logger, c *cli.Context) (clickhouse.Conn, func(context.Context) error, error) { logger = logger.With(attr.SlogComponent("clickhouse")) nilFunc := noopShutdown @@ -203,6 +209,7 @@ func newClickhouseClient(ctx context.Context, logger *slog.Logger, c *cli.Contex MaxOpenConns: 32, MaxIdleConns: 16, DialTimeout: 10 * time.Second, + ReadTimeout: clickhouseSocketReadTimeout, TLS: &tls.Config{ // #nosec G402 -- we're reading the value from an environment variable. InsecureSkipVerify: insecure, @@ -260,7 +267,14 @@ func newClickhouseClient(ctx context.Context, logger *slog.Logger, c *cli.Contex // forwards the caller's span context to ClickHouse by default, so // server-side spans (system.opentelemetry_span_log) can be joined against // APM traces by trace id; no per-call-site wiring exists or is needed. - return o11y.TraceClickhouseConn(conn), shutdown, nil + resilientConn := clickhouseclient.WithReadResilience( + conn, + func() (clickhouse.Conn, error) { + return clickhouse.Open(opts) + }, + clickhouseReadOperationTimeout, + ) + return o11y.TraceClickhouseConn(resilientConn), shutdown, nil } type dbClientOptions struct { diff --git a/server/internal/clickhouseclient/read.go b/server/internal/clickhouseclient/read.go new file mode 100644 index 00000000000..471e855490d --- /dev/null +++ b/server/internal/clickhouseclient/read.go @@ -0,0 +1,222 @@ +package clickhouseclient + +import ( + "context" + "errors" + "io" + "net" + "sync" + "time" + + "github.com/ClickHouse/clickhouse-go/v2" + "github.com/ClickHouse/clickhouse-go/v2/lib/driver" +) + +type readResilientConn struct { + clickhouse.Conn + newConn func() (clickhouse.Conn, error) + timeout time.Duration +} + +// WithReadResilience bounds read operations and retries transport failures that +// happen before a result is exposed. The retry uses a newly opened connection; +// writes and partially consumed result sets are never retried. +func WithReadResilience( + conn clickhouse.Conn, + newConn func() (clickhouse.Conn, error), + timeout time.Duration, +) clickhouse.Conn { + if conn == nil { + panic("clickhouse connection is nil") + } + if newConn == nil { + panic("clickhouse connection factory is nil") + } + if timeout <= 0 { + panic("clickhouse read timeout must be positive") + } + + return &readResilientConn{ + Conn: conn, + newConn: newConn, + timeout: timeout, + } +} + +//nolint:wrapcheck // Preserve ClickHouse errors so callers can classify them. +func (c *readResilientConn) Select(ctx context.Context, dest any, query string, args ...any) error { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + defer cancel() + + // Select can mutate dest before returning an error, so it is bounded but not + // retried. Callers that need retry semantics should use Query or QueryRow. + return c.Conn.Select(ctx, dest, query, args...) +} + +//nolint:wrapcheck // Preserve ClickHouse errors so callers can classify them. +func (c *readResilientConn) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + + rows, err := c.Conn.Query(ctx, query, args...) + if !isRetryableReadError(ctx, err) { + if err != nil { + cancel() + return nil, err + } + return &managedRows{ + Rows: rows, + cancel: cancel, + conn: nil, + once: sync.Once{}, + err: nil, + }, nil + } + + fresh, openErr := c.newConn() + if openErr != nil { + cancel() + return nil, errors.Join(err, openErr) + } + + retryRows, retryErr := fresh.Query(ctx, query, args...) + if retryErr != nil { + cancel() + return nil, errors.Join(err, retryErr, fresh.Close()) + } + + return &managedRows{ + Rows: retryRows, + cancel: cancel, + conn: fresh, + once: sync.Once{}, + err: nil, + }, nil +} + +func (c *readResilientConn) QueryRow(ctx context.Context, query string, args ...any) driver.Row { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + + row := c.Conn.QueryRow(ctx, query, args...) + initialErr := row.Err() + if !isRetryableReadError(ctx, initialErr) { + return &managedRow{ + Row: row, + cancel: cancel, + conn: nil, + initialErr: nil, + once: sync.Once{}, + err: nil, + } + } + + fresh, openErr := c.newConn() + if openErr != nil { + cancel() + return errorRow{err: errors.Join(initialErr, openErr)} + } + + retryRow := fresh.QueryRow(ctx, query, args...) + return &managedRow{ + Row: retryRow, + cancel: cancel, + conn: fresh, + initialErr: initialErr, + once: sync.Once{}, + err: nil, + } +} + +func isRetryableReadError(ctx context.Context, err error) bool { + if err == nil || ctx.Err() != nil { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, net.ErrClosed) { + return true + } + + var netErr net.Error + return errors.As(err, &netErr) +} + +type managedRows struct { + driver.Rows + cancel context.CancelFunc + conn clickhouse.Conn + once sync.Once + err error +} + +func (r *managedRows) Next() bool { + if r.Rows.Next() { + return true + } + + _ = r.cleanup() + return false +} + +func (r *managedRows) Close() error { + return errors.Join(r.Rows.Close(), r.cleanup()) +} + +func (r *managedRows) cleanup() error { + r.once.Do(func() { + r.cancel() + if r.conn != nil { + r.err = r.conn.Close() + } + }) + return r.err +} + +type managedRow struct { + driver.Row + cancel context.CancelFunc + conn clickhouse.Conn + initialErr error + once sync.Once + err error +} + +func (r *managedRow) Err() error { + err := r.Row.Err() + if err == nil { + return nil + } + + return errors.Join(r.initialErr, err, r.cleanup()) +} + +func (r *managedRow) Scan(dest ...any) error { + err := r.Row.Scan(dest...) + if err != nil { + return errors.Join(r.initialErr, err, r.cleanup()) + } + return r.cleanup() +} + +func (r *managedRow) ScanStruct(dest any) error { + err := r.Row.ScanStruct(dest) + if err != nil { + return errors.Join(r.initialErr, err, r.cleanup()) + } + return r.cleanup() +} + +func (r *managedRow) cleanup() error { + r.once.Do(func() { + r.cancel() + if r.conn != nil { + r.err = r.conn.Close() + } + }) + return r.err +} + +type errorRow struct { + err error +} + +func (r errorRow) Err() error { return r.err } +func (r errorRow) Scan(...any) error { return r.err } +func (r errorRow) ScanStruct(any) error { return r.err } diff --git a/server/internal/clickhouseclient/read_test.go b/server/internal/clickhouseclient/read_test.go new file mode 100644 index 00000000000..d8029c6f793 --- /dev/null +++ b/server/internal/clickhouseclient/read_test.go @@ -0,0 +1,248 @@ +package clickhouseclient + +import ( + "context" + "errors" + "fmt" + "io" + "testing" + "time" + + "github.com/ClickHouse/clickhouse-go/v2" + "github.com/ClickHouse/clickhouse-go/v2/lib/driver" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +func TestWithReadResilienceRetriesQueryOnFreshConnection(t *testing.T) { + t.Parallel() + + primary := new(mockConn) + fresh := new(mockConn) + rows := new(mockRows) + firstErr := fmt.Errorf("read first block: %w", io.EOF) + var firstCtx context.Context + var retryCtx context.Context + + primary.On("Query", mock.Anything, "SELECT 1", mock.Anything). + Run(func(args mock.Arguments) { + ctx, ok := args.Get(0).(context.Context) + require.True(t, ok) + firstCtx = ctx + }). + Return(nil, firstErr). + Once() + fresh.On("Query", mock.Anything, "SELECT 1", mock.Anything). + Run(func(args mock.Arguments) { + ctx, ok := args.Get(0).(context.Context) + require.True(t, ok) + retryCtx = ctx + }). + Return(rows, nil). + Once() + rows.On("Next").Return(false).Once() + rows.On("Close").Return(nil).Once() + fresh.On("Close").Return(nil).Once() + + factoryCalls := 0 + conn := WithReadResilience(primary, func() (clickhouse.Conn, error) { + factoryCalls++ + return fresh, nil + }, time.Minute) + + gotRows, err := conn.Query(t.Context(), "SELECT 1") + require.NoError(t, err) + require.NotNil(t, gotRows) + require.Equal(t, 1, factoryCalls) + + firstDeadline, firstHasDeadline := firstCtx.Deadline() + retryDeadline, retryHasDeadline := retryCtx.Deadline() + require.True(t, firstHasDeadline) + require.True(t, retryHasDeadline) + require.Equal(t, firstDeadline, retryDeadline) + + require.False(t, gotRows.Next()) + require.ErrorIs(t, firstCtx.Err(), context.Canceled) + require.NoError(t, gotRows.Close()) + primary.AssertExpectations(t) + fresh.AssertExpectations(t) + rows.AssertExpectations(t) +} + +func TestWithReadResilienceBoundsQueryAndDoesNotRetryDeadline(t *testing.T) { + t.Parallel() + + primary := new(mockConn) + primary.On("Query", mock.Anything, "SELECT sleep", mock.Anything). + Run(func(args mock.Arguments) { + ctx, ok := args.Get(0).(context.Context) + require.True(t, ok) + <-ctx.Done() + }). + Return(nil, context.DeadlineExceeded). + Once() + + factoryCalls := 0 + conn := WithReadResilience(primary, func() (clickhouse.Conn, error) { + factoryCalls++ + return nil, errors.New("must not open a retry connection") + }, time.Nanosecond) + + rows, err := conn.Query(t.Context(), "SELECT sleep") + require.Nil(t, rows) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Zero(t, factoryCalls) + primary.AssertExpectations(t) +} + +func TestWithReadResilienceDoesNotRetryServerError(t *testing.T) { + t.Parallel() + + primary := new(mockConn) + serverErr := errors.New("clickhouse exception: invalid query") + primary.On("Query", mock.Anything, "SELECT invalid", mock.Anything). + Return(nil, serverErr). + Once() + + factoryCalls := 0 + conn := WithReadResilience(primary, func() (clickhouse.Conn, error) { + factoryCalls++ + return nil, errors.New("must not open a retry connection") + }, time.Minute) + + rows, err := conn.Query(t.Context(), "SELECT invalid") + require.Nil(t, rows) + require.ErrorIs(t, err, serverErr) + require.Zero(t, factoryCalls) + primary.AssertExpectations(t) +} + +func TestWithReadResilienceRetriesQueryRowOnFreshConnection(t *testing.T) { + t.Parallel() + + primary := new(mockConn) + fresh := new(mockConn) + initialRow := new(mockRow) + retryRow := new(mockRow) + firstErr := fmt.Errorf("read first block: %w", io.EOF) + var retryCtx context.Context + + primary.On("QueryRow", mock.Anything, "SELECT count()", mock.Anything). + Return(initialRow). + Once() + initialRow.On("Err").Return(firstErr).Once() + fresh.On("QueryRow", mock.Anything, "SELECT count()", mock.Anything). + Run(func(args mock.Arguments) { + ctx, ok := args.Get(0).(context.Context) + require.True(t, ok) + retryCtx = ctx + }). + Return(retryRow). + Once() + retryRow.On("Scan", mock.Anything). + Run(func(args mock.Arguments) { + dest, ok := args.Get(0).([]any) + require.True(t, ok) + require.Len(t, dest, 1) + count, ok := dest[0].(*uint64) + require.True(t, ok) + *count = 42 + }). + Return(nil). + Once() + fresh.On("Close").Return(nil).Once() + + conn := WithReadResilience(primary, func() (clickhouse.Conn, error) { + return fresh, nil + }, time.Minute) + + var count uint64 + err := conn.QueryRow(t.Context(), "SELECT count()").Scan(&count) + require.NoError(t, err) + require.Equal(t, uint64(42), count) + require.ErrorIs(t, retryCtx.Err(), context.Canceled) + primary.AssertExpectations(t) + fresh.AssertExpectations(t) + initialRow.AssertExpectations(t) + retryRow.AssertExpectations(t) +} + +func TestWithReadResilienceDoesNotRetryWrites(t *testing.T) { + t.Parallel() + + primary := new(mockConn) + primary.On("Exec", mock.Anything, "INSERT INTO events VALUES (?)", mock.Anything). + Return(nil). + Once() + + factoryCalls := 0 + conn := WithReadResilience(primary, func() (clickhouse.Conn, error) { + factoryCalls++ + return nil, errors.New("must not open a retry connection") + }, time.Minute) + + require.NoError(t, conn.Exec(t.Context(), "INSERT INTO events VALUES (?)", 42)) + require.Zero(t, factoryCalls) + primary.AssertExpectations(t) +} + +type mockConn struct { + mock.Mock + clickhouse.Conn +} + +func (m *mockConn) Select(ctx context.Context, dest any, query string, args ...any) error { + return m.Called(ctx, dest, query, args).Error(0) +} + +func (m *mockConn) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) { + call := m.Called(ctx, query, args) + rows, _ := call.Get(0).(driver.Rows) + return rows, call.Error(1) +} + +func (m *mockConn) QueryRow(ctx context.Context, query string, args ...any) driver.Row { + row, ok := m.Called(ctx, query, args).Get(0).(driver.Row) + if !ok { + panic("mock QueryRow result is not driver.Row") + } + return row +} + +func (m *mockConn) Exec(ctx context.Context, query string, args ...any) error { + return m.Called(ctx, query, args).Error(0) +} + +func (m *mockConn) Close() error { + return m.Called().Error(0) +} + +type mockRows struct { + mock.Mock + driver.Rows +} + +func (m *mockRows) Next() bool { + return m.Called().Bool(0) +} + +func (m *mockRows) Close() error { + return m.Called().Error(0) +} + +type mockRow struct { + mock.Mock + driver.Row +} + +func (m *mockRow) Err() error { + return m.Called().Error(0) +} + +func (m *mockRow) Scan(dest ...any) error { + return m.Called(dest).Error(0) +} + +func (m *mockRow) ScanStruct(dest any) error { + return m.Called(dest).Error(0) +}