Decide in the Postgres writer whether a failed chunk may have been written, drop the dialect hook

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])
 	}

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)

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)

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
 }
 

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")))
+}

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
-}