From fdc88a3163278b1cd2e79adb8765883ac4b789dc Mon Sep 17 00:00:00 2001 From: binwiederhier Date: Sat, 31 Oct 2026 10:34:52 +0000 Subject: [PATCH] Decide in the Postgres writer whether a failed chunk may have been written, drop the dialect hook --- limit/service.go | 3 --- limit/service_test.go | 2 +- limit/store.go | 22 ++++++++++++++++------ limit/store_postgres.go | 8 ++++++-- limit/store_postgres_test.go | 19 +++++++++++++++++++ limit/store_sqlite.go | 7 ------- 6 files changed, 42 insertions(+), 19 deletions(-) create mode 100644 limit/store_postgres_test.go diff --git a/limit/service.go b/limit/service.go index 8653e0e1..24ab85f6 100644 --- a/limit/service.go +++ b/limit/service.go @@ -401,9 +401,6 @@ func (s *Service) write(ctx context.Context) { return } written, err := s.store.write(ctx, rows, s.now().Unix()) - if err != nil && s.store.mayHaveCommitted(err) { - written = min(written+writeChunkRows, len(rows)) - } for i := range rows[:written] { s.retire(&rows[i]) } diff --git a/limit/service_test.go b/limit/service_test.go index 229180f6..4250d455 100644 --- a/limit/service_test.go +++ b/limit/service_test.go @@ -822,7 +822,7 @@ func TestService_Persist_TimeoutAfterCommitIsNotCountedTwice(t *testing.T) { if err := writeChunk(ctx, rows, now); err != nil { return err } - return context.DeadlineExceeded // Committed, but the answer came too late + return fmt.Errorf("%w: %w", errMaybeWritten, context.DeadlineExceeded) // Committed, but the answer came too late } s := startTestService(t, store, time.Now) require.True(t, s.Allow("user:u_1", Rule{Kind: KindMessages}, 26).OK) diff --git a/limit/store.go b/limit/store.go index 5a006cc6..5e2804b4 100644 --- a/limit/store.go +++ b/limit/store.go @@ -2,6 +2,7 @@ package limit import ( "context" + "errors" "slices" "strings" @@ -17,6 +18,12 @@ const ( pruneChunkRows = 10000 ) +var ( + // errMaybeWritten wraps a chunk's error when the chunk may have been written anyway (a timeout + // after the statement was sent), so the caller counts it as written rather than writing it twice + errMaybeWritten = errors.New("chunk may have been written") +) + // usageRow is one key's daily counts in one period: a delta on the way to the store, the row's // totals on the way back type usageRow struct { @@ -29,10 +36,9 @@ type usageRow struct { // stamped (updated_at) with the writing node's clock, like the reader's watermark, so the // database's clock does not matter. Only the service's sync worker calls it. type Store struct { - db *db.DB - queries *queries - writeChunk func(ctx context.Context, rows []usageRow, now int64) error // Dialect specific: unnest on Postgres, a prepared statement per row on SQLite - mayHaveCommitted func(err error) bool // Dialect specific: whether a failed chunk may have been written anyway + db *db.DB + queries *queries + writeChunk func(ctx context.Context, rows []usageRow, now int64) error // Dialect specific: unnest on Postgres, a prepared statement per row on SQLite } // queries are the dialect's SQL statements @@ -46,8 +52,9 @@ type queries struct { // row's period is older. Rows of an older period than the stored one are added to the stored // one (a node's leftover of a finished period, see the usage queries). Rows are written sorted // by key in chunks, so two nodes lock overlapping keys in the same order; it returns how many -// rows (in sorted order) are committed: on an error, the chunks before it are. The rows' -// updated_at is now (unix seconds). +// rows (in sorted order) are committed: on an error, the chunks before it are, plus the failed +// chunk if it may have been written anyway (errMaybeWritten). The rows' updated_at is now (unix +// seconds). func (s *Store) write(ctx context.Context, rows []usageRow, now int64) (int, error) { slices.SortFunc(rows, func(a, b usageRow) int { return strings.Compare(a.key, b.key) @@ -55,6 +62,9 @@ func (s *Store) write(ctx context.Context, rows []usageRow, now int64) (int, err written := 0 for chunk := range slices.Chunk(rows, writeChunkRows) { if err := s.writeChunk(ctx, chunk, now); err != nil { + if errors.Is(err, errMaybeWritten) { + written += len(chunk) + } return written, err } written += len(chunk) diff --git a/limit/store_postgres.go b/limit/store_postgres.go index 19f11961..b8154138 100644 --- a/limit/store_postgres.go +++ b/limit/store_postgres.go @@ -3,6 +3,7 @@ package limit import ( "context" "errors" + "fmt" "github.com/jackc/pgx/v5/pgconn" @@ -39,11 +40,11 @@ func NewPostgresStore(d *db.DB) *Store { }, } s.writeChunk = s.writePostgres - s.mayHaveCommitted = mayHaveCommittedPostgres return s } -// writePostgres writes the rows in one statement, as arrays +// writePostgres writes the rows in one statement, as arrays. A timeout after the statement was +// sent is wrapped in errMaybeWritten, since the commit may have been acknowledged after the deadline. func (s *Store) writePostgres(ctx context.Context, rows []usageRow, now int64) error { keys := make([]string, len(rows)) periodStarts := make([]int64, len(rows)) @@ -58,6 +59,9 @@ func (s *Store) writePostgres(ctx context.Context, rows []usageRow, now int64) e } } _, err := s.db.ExecContext(ctx, s.queries.upsert, keys, periodStarts, counts[0], counts[1], counts[2], now) + if err != nil && mayHaveCommittedPostgres(err) { + return fmt.Errorf("%w: %w", errMaybeWritten, err) + } return err } diff --git a/limit/store_postgres_test.go b/limit/store_postgres_test.go new file mode 100644 index 00000000..036ac0d7 --- /dev/null +++ b/limit/store_postgres_test.go @@ -0,0 +1,19 @@ +package limit + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/require" +) + +func TestStore_MayHaveCommittedPostgres(t *testing.T) { + // A timeout may have committed (the answer came after the deadline); a server error or an + // error before sending never did. A timeout waiting for a pool connection is not told apart + // from one after sending, so it counts as written too (counts may be lost, never doubled). + require.True(t, mayHaveCommittedPostgres(context.DeadlineExceeded)) + require.False(t, mayHaveCommittedPostgres(&pgconn.PgError{Code: "40001"})) + require.False(t, mayHaveCommittedPostgres(errors.New("connection refused"))) +} diff --git a/limit/store_sqlite.go b/limit/store_sqlite.go index 5e33de54..dd9b1d44 100644 --- a/limit/store_sqlite.go +++ b/limit/store_sqlite.go @@ -36,7 +36,6 @@ func NewSQLiteStore(d *db.DB) *Store { }, } s.writeChunk = s.writeSQLite - s.mayHaveCommitted = mayHaveCommittedSQLite return s } @@ -56,9 +55,3 @@ func (s *Store) writeSQLite(ctx context.Context, rows []usageRow, now int64) err return nil }) } - -// mayHaveCommittedSQLite reports false: a chunk is written in one local transaction, which is -// rolled back on any error, so a failed chunk was never written -func mayHaveCommittedSQLite(_ error) bool { - return false -}