From 345f1d2d1ec9c3379751f62753c09c492100c97a Mon Sep 17 00:00:00 2001 From: atharrva01 Date: Mon, 10 Aug 2026 06:25:47 +0530 Subject: [PATCH] fix(recovery): persist audit recovery claims postgres.AuditTransactionStore only overrode WriteDB, GetSchema and CreateSchema, so it inherited ClaimPendingTransactions and ReleaseRecoveryClaim from sqlcommon. The common implementations are a plain SELECT and a no-op, neither of which persists a claim, so every replica selected the same pending audit transactions on every tick and processed all of them. The audit store now uses the same atomic UPDATE ... RETURNING claim the owner store has had since it was introduced. The SQL is identical for the two stores and only the requests table differs, so it moves to a shared recoveryClaimStore that both hold as a field. It is a named field rather than an embedded one because sqlcommon.TransactionStore also provides these methods, and embedding both at the same depth makes the selectors ambiguous, which silently drops them from the method set. No schema change. recovery_claimed_by, recovery_claim_expires_at and the two supporting indexes already come from the shared sqlcommon schema that audit storage uses. CleanupExpiredClaims is not exposed on the audit store. Nothing calls it on either path and the claim query already reclaims expired leases inline. Signed-off-by: atharrva01 --- docs/services/storage/recovery.md | 13 + .../sql/postgres/audit_recovery_claim_test.go | 261 ++++++++++++++++++ .../db/sql/postgres/recovery_claim_test.go | 23 +- .../storage/db/sql/postgres/transactions.go | 73 ++++- 4 files changed, 351 insertions(+), 19 deletions(-) create mode 100644 token/services/storage/db/sql/postgres/audit_recovery_claim_test.go diff --git a/docs/services/storage/recovery.md b/docs/services/storage/recovery.md index 964fc2045f..d29243462a 100644 --- a/docs/services/storage/recovery.md +++ b/docs/services/storage/recovery.md @@ -48,6 +48,19 @@ PostgreSQL is the recommended database for production multi-instance deployments - Supports horizontal scaling with multiple replicas - Leader election prevents conflicting recovery attempts +A recovery manager is started per TMS for both owner and audit storage. Both stores use the +same atomic claim: a pending row is handed to exactly one replica for the duration of its +lease, and `ReleaseRecoveryClaim` frees it again as soon as the sweep is done with it. Claims +live in the `recovery_claimed_by` and `recovery_claim_expires_at` columns of each store's own +requests table, so an audit claim never hides a row from the owner sweep or the other way +round. + +Leader election currently differs between the two. The owner store elects a leader through a +PostgreSQL advisory lock; the audit store is built without a leader factory, so +`AcquireRecoveryLeadership` grants leadership locally to every replica. Audit sweeps therefore +run everywhere at once, and it is the atomic claim rather than leader election that keeps each +pending audit transaction from being processed more than once. + ### SQLite (Development and Single-Node) SQLite is supported for single-node deployments and development: diff --git a/token/services/storage/db/sql/postgres/audit_recovery_claim_test.go b/token/services/storage/db/sql/postgres/audit_recovery_claim_test.go new file mode 100644 index 0000000000..8ba20ea02d --- /dev/null +++ b/token/services/storage/db/sql/postgres/audit_recovery_claim_test.go @@ -0,0 +1,261 @@ +/* +Copyright IBM Corp. All Rights Reserved. + +SPDX-License-Identifier: Apache-2.0 +*/ + +package postgres + +import ( + "context" + "math/big" + "testing" + "time" + + tokensdriver "github.com/LFDT-Panurus/panurus/token/services/storage/db/driver" + "github.com/stretchr/testify/require" +) + +// newAuditStore returns an AuditTransactionStore backed by the given connection string, +// with its schema already created. +func newAuditStore(t *testing.T, pgConnStr, name string) *AuditTransactionStore { + t.Helper() + + storeInterface, err := NewDriver(postgresCfg(pgConnStr, name)).NewAuditTransaction("test", name) + require.NoError(t, err) + store, ok := storeInterface.(*AuditTransactionStore) + require.True(t, ok) + require.NoError(t, store.CreateSchema()) + + return store +} + +// addPendingAuditTx inserts a Pending audit transaction and backdates it so it falls +// inside the recovery claim window. +func addPendingAuditTx(t *testing.T, ctx context.Context, store *AuditTransactionStore, storedAt time.Time, txIDs ...string) { + t.Helper() + + aw, err := store.NewTransactionStoreTransaction() + require.NoError(t, err) + + for _, txID := range txIDs { + require.NoError(t, aw.AddTokenRequest(ctx, txID, []byte("request"), nil, nil, []byte("hash"))) + require.NoError(t, aw.AddTransaction(ctx, tokensdriver.TransactionRecord{ + TxID: txID, + ActionType: tokensdriver.Transfer, + SenderEID: "sender", + RecipientEID: "recipient", + TokenType: "USD", + Amount: big.NewInt(100), + Timestamp: storedAt, + })) + } + require.NoError(t, aw.Commit()) + + ageRequests(t, ctx, store.claims, storedAt, txIDs...) +} + +// TestAuditClaimPendingTransactions_Atomic is the audit-side counterpart of +// TestClaimPendingTransactions_Atomic. Before the audit store overrode +// ClaimPendingTransactions it inherited the permissive SELECT from sqlcommon, which +// persisted no claim: both replicas below would have received all five transactions and +// each would have processed the whole batch on every sweep. +func TestAuditClaimPendingTransactions_Atomic(t *testing.T) { + terminate, pgConnStr := startContainer(t) + defer terminate() + + ctx := context.Background() + + // Two stores over the same database, standing in for two replicas. + store1 := newAuditStore(t, pgConnStr, "aud_atomic") + store2 := newAuditStore(t, pgConnStr, "aud_atomic") + + now := time.Now().UTC() + oldTime := now.Add(-10 * time.Minute) + + txIDs := []string{"atx1", "atx2", "atx3", "atx4", "atx5"} + addPendingAuditTx(t, ctx, store1, oldTime, txIDs...) + + params := tokensdriver.RecoveryClaimParams{ + OlderThan: now, + LeaseDuration: 5 * time.Minute, + Limit: 10, + Owner: "instance1", + } + + claimed1, err := store1.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, claimed1, len(txIDs), "first replica should claim every pending audit transaction") + + params.Owner = "instance2" + claimed2, err := store2.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Empty(t, claimed2, "second replica must claim nothing, the rows are already claimed") +} + +// TestAuditClaimPendingTransactions_Lease verifies an audit claim is exclusive until its +// lease expires, and reclaimable afterwards. +func TestAuditClaimPendingTransactions_Lease(t *testing.T) { + terminate, pgConnStr := startContainer(t) + defer terminate() + + ctx := context.Background() + store := newAuditStore(t, pgConnStr, "aud_lease") + + now := time.Now().UTC() + oldTime := now.Add(-10 * time.Minute) + addPendingAuditTx(t, ctx, store, oldTime, "atx1") + + params := tokensdriver.RecoveryClaimParams{ + OlderThan: now, + LeaseDuration: 1 * time.Second, + Limit: 10, + Owner: "instance1", + } + + claimed, err := store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, claimed, 1) + + params.Owner = "instance2" + claimed, err = store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Empty(t, claimed, "another owner must not claim before the lease expires") + + require.Eventually(t, func() bool { + claimed, err = store.ClaimPendingTransactions(ctx, params) + + return err == nil && len(claimed) == 1 + }, 10*time.Second, 250*time.Millisecond, "claim should become available once the lease expires") +} + +// TestAuditReleaseRecoveryClaim verifies the audit store actually clears the claim, so the +// row is immediately available again rather than staying locked until the lease runs out. +// The inherited implementation is a no-op and would leave the claim in place. +func TestAuditReleaseRecoveryClaim(t *testing.T) { + terminate, pgConnStr := startContainer(t) + defer terminate() + + ctx := context.Background() + store := newAuditStore(t, pgConnStr, "aud_release") + + now := time.Now().UTC() + oldTime := now.Add(-10 * time.Minute) + addPendingAuditTx(t, ctx, store, oldTime, "atx1") + + params := tokensdriver.RecoveryClaimParams{ + OlderThan: now, + LeaseDuration: 5 * time.Minute, + Limit: 10, + Owner: "instance1", + } + + claimed, err := store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, claimed, 1) + + require.NoError(t, store.ReleaseRecoveryClaim(ctx, "atx1", "instance1", "recovered successfully")) + + // A different owner can claim it right away, well inside the original 5 minute lease. + params.Owner = "instance2" + claimed, err = store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, claimed, 1, "released claim should be immediately available to another owner") + + status, message, err := store.GetStatus(ctx, "atx1") + require.NoError(t, err) + require.Equal(t, tokensdriver.Pending, status) + require.Equal(t, "recovered successfully", message) +} + +// TestAuditReleaseRecoveryClaim_WrongOwner verifies a replica cannot release a claim held +// by another replica. +func TestAuditReleaseRecoveryClaim_WrongOwner(t *testing.T) { + terminate, pgConnStr := startContainer(t) + defer terminate() + + ctx := context.Background() + store := newAuditStore(t, pgConnStr, "aud_relowner") + + now := time.Now().UTC() + oldTime := now.Add(-10 * time.Minute) + addPendingAuditTx(t, ctx, store, oldTime, "atx1") + + params := tokensdriver.RecoveryClaimParams{ + OlderThan: now, + LeaseDuration: 5 * time.Minute, + Limit: 10, + Owner: "instance1", + } + + claimed, err := store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, claimed, 1) + + // Releasing under the wrong owner is not an error, it simply affects no rows. + require.NoError(t, store.ReleaseRecoveryClaim(ctx, "atx1", "instance2", "should not apply")) + + params.Owner = "instance3" + claimed, err = store.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Empty(t, claimed, "claim held by instance1 must survive a release attempt by instance2") +} + +// TestAuditAndOwnerClaimsAreIndependent pins that the two stores keep their claims in +// their own requests tables. A claim taken on the audit side must not hide a pending row +// from the owner sweep, or vice versa. +func TestAuditAndOwnerClaimsAreIndependent(t *testing.T) { + terminate, pgConnStr := startContainer(t) + defer terminate() + + ctx := context.Background() + + auditStore := newAuditStore(t, pgConnStr, "aud_split") + + ownerInterface, err := NewDriver(postgresCfg(pgConnStr, "aud_split")). + NewOwnerTransaction("test", "aud_split") + require.NoError(t, err) + ownerStore, ok := ownerInterface.(*TransactionStore) + require.True(t, ok) + require.NoError(t, ownerStore.CreateSchema()) + + require.NotEqual(t, auditStore.claims.tables.Requests, ownerStore.claims.tables.Requests, + "audit and owner stores must not share a requests table") + + now := time.Now().UTC() + oldTime := now.Add(-10 * time.Minute) + + addPendingAuditTx(t, ctx, auditStore, oldTime, "atx1") + + aw, err := ownerStore.NewTransactionStoreTransaction() + require.NoError(t, err) + require.NoError(t, aw.AddTokenRequest(ctx, "otx1", []byte("request"), nil, nil, []byte("hash"))) + require.NoError(t, aw.AddTransaction(ctx, tokensdriver.TransactionRecord{ + TxID: "otx1", + ActionType: tokensdriver.Transfer, + SenderEID: "sender", + RecipientEID: "recipient", + TokenType: "USD", + Amount: big.NewInt(100), + Timestamp: oldTime, + })) + require.NoError(t, aw.Commit()) + ageRequests(t, ctx, ownerStore.claims, oldTime, "otx1") + + params := tokensdriver.RecoveryClaimParams{ + OlderThan: now, + LeaseDuration: 5 * time.Minute, + Limit: 10, + Owner: "instance1", + } + + auditClaimed, err := auditStore.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, auditClaimed, 1) + require.Equal(t, "atx1", auditClaimed[0].TxID) + + ownerClaimed, err := ownerStore.ClaimPendingTransactions(ctx, params) + require.NoError(t, err) + require.Len(t, ownerClaimed, 1, "owner sweep must be unaffected by the audit claim") + require.Equal(t, "otx1", ownerClaimed[0].TxID) +} diff --git a/token/services/storage/db/sql/postgres/recovery_claim_test.go b/token/services/storage/db/sql/postgres/recovery_claim_test.go index 37df594246..ad4f2e5f31 100644 --- a/token/services/storage/db/sql/postgres/recovery_claim_test.go +++ b/token/services/storage/db/sql/postgres/recovery_claim_test.go @@ -71,7 +71,7 @@ func TestClaimPendingTransactions_Atomic(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store1, oldTime, txIDs...) + ageRequests(t, ctx, store1.claims, oldTime, txIDs...) // Both instances try to claim the same transactions params := tokensdriver.RecoveryClaimParams{ @@ -132,7 +132,7 @@ func TestClaimPendingTransactions_Lease(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store, oldTime, txID) + ageRequests(t, ctx, store.claims, oldTime, txID) // Claim with very short lease params := tokensdriver.RecoveryClaimParams{ @@ -201,7 +201,7 @@ func TestClaimPendingTransactions_Idempotent(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store, oldTime, txID) + ageRequests(t, ctx, store.claims, oldTime, txID) // Claim transaction params := tokensdriver.RecoveryClaimParams{ @@ -266,7 +266,7 @@ func TestClaimPendingTransactions_Limit(t *testing.T) { err = aw.Commit() require.NoError(t, err) for i, txID := range txIDs { - ageRequests(t, ctx, store, oldTime.Add(time.Duration(i)*time.Second), txID) + ageRequests(t, ctx, store.claims, oldTime.Add(time.Duration(i)*time.Second), txID) } // Claim with limit of 3 @@ -328,7 +328,7 @@ func TestReleaseRecoveryClaim(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store, oldTime, txID) + ageRequests(t, ctx, store.claims, oldTime, txID) // Claim transaction params := tokensdriver.RecoveryClaimParams{ @@ -393,7 +393,7 @@ func TestReleaseRecoveryClaim_WrongOwner(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store, oldTime, txID) + ageRequests(t, ctx, store.claims, oldTime, txID) // Claim transaction params := tokensdriver.RecoveryClaimParams{ @@ -461,7 +461,7 @@ func TestCleanupExpiredClaims(t *testing.T) { err = aw.Commit() require.NoError(t, err) - ageRequests(t, ctx, store, oldTime, txIDs...) + ageRequests(t, ctx, store.claims, oldTime, txIDs...) // Claim with very short lease params := tokensdriver.RecoveryClaimParams{ @@ -490,13 +490,16 @@ func TestCleanupExpiredClaims(t *testing.T) { require.Len(t, claimed, 3, "Should be able to claim after cleanup") } -func ageRequests(t *testing.T, ctx context.Context, store *TransactionStore, storedAt time.Time, txIDs ...string) { +// ageRequests backdates stored_at so the rows fall inside the recovery claim window. +// It takes the claim store rather than a concrete store type so the owner and audit +// tests can share it. +func ageRequests(t *testing.T, ctx context.Context, claims *recoveryClaimStore, storedAt time.Time, txIDs ...string) { t.Helper() // #nosec G201 -- table name comes from the test-created store. - query := fmt.Sprintf("UPDATE %s SET stored_at = $1 WHERE tx_id = $2", store.tables.Requests) + query := fmt.Sprintf("UPDATE %s SET stored_at = $1 WHERE tx_id = $2", claims.tables.Requests) for _, txID := range txIDs { - result, err := store.writeDB.ExecContext(ctx, query, storedAt, txID) + result, err := claims.writeDB.ExecContext(ctx, query, storedAt, txID) require.NoError(t, err) rowsAffected, err := result.RowsAffected() diff --git a/token/services/storage/db/sql/postgres/transactions.go b/token/services/storage/db/sql/postgres/transactions.go index 78cc92e8b3..94d0458596 100644 --- a/token/services/storage/db/sql/postgres/transactions.go +++ b/token/services/storage/db/sql/postgres/transactions.go @@ -28,6 +28,7 @@ type AuditTransactionStore struct { *sqlcommon.TransactionStore writeDB *sql.DB lockID int64 + claims *recoveryClaimStore } // WriteDB returns the underlying write *sql.DB. @@ -49,10 +50,64 @@ func (s *AuditTransactionStore) CreateSchema() error { // TransactionStore extends the common TransactionStore with PostgreSQL-specific atomic claim operations. type TransactionStore struct { *sqlcommon.TransactionStore + writeDB *sql.DB + lockID int64 + claims *recoveryClaimStore +} + +// recoveryClaimStore holds the PostgreSQL-specific recovery claim operations shared by the +// owner and audit transaction stores. Both stores keep their claim state in the same +// columns of their own requests table, so the SQL is identical and only the table differs. +// +// It is a named field rather than an embedded one on purpose: sqlcommon.TransactionStore +// already provides ClaimPendingTransactions and ReleaseRecoveryClaim, and embedding both at +// the same depth would make those selectors ambiguous. Go then drops them from the method +// set and the store silently stops satisfying driver.TransactionStore. Explicit forwarding +// keeps the override visible at the call site. +type recoveryClaimStore struct { readDB *sql.DB writeDB *sql.DB tables sqlcommon.TableNames - lockID int64 +} + +func newRecoveryClaimStore(dbs *scommon.RWDB, tableNames sqlcommon.TableNames) *recoveryClaimStore { + return &recoveryClaimStore{ + readDB: dbs.ReadDB, + writeDB: dbs.WriteDB, + tables: tableNames, + } +} + +// ClaimPendingTransactions atomically claims a batch of pending transactions. +func (db *TransactionStore) ClaimPendingTransactions(ctx context.Context, params tokensdriver.RecoveryClaimParams) ([]*tokensdriver.RecoveryClaim, error) { + return db.claims.claimPending(ctx, params) +} + +// ReleaseRecoveryClaim releases the recovery claim on a transaction. +func (db *TransactionStore) ReleaseRecoveryClaim(ctx context.Context, txID string, owner string, message string) error { + return db.claims.releaseClaim(ctx, txID, owner, message) +} + +// CleanupExpiredClaims removes expired recovery claims. Returns the number of claims cleaned up. +func (db *TransactionStore) CleanupExpiredClaims(ctx context.Context) (int, error) { + return db.claims.cleanupExpired(ctx) +} + +// ClaimPendingTransactions atomically claims a batch of pending audit transactions. +// +// Without this override the store inherits the permissive SELECT from +// sqlcommon.TransactionStore, which persists no claim at all: every replica selects the +// same pending rows on every sweep and processes all of them. The owner path has had the +// atomic claim since it was introduced; audit was simply never given it. +func (s *AuditTransactionStore) ClaimPendingTransactions(ctx context.Context, params tokensdriver.RecoveryClaimParams) ([]*tokensdriver.RecoveryClaim, error) { + return s.claims.claimPending(ctx, params) +} + +// ReleaseRecoveryClaim releases the recovery claim on an audit transaction. +// The inherited implementation is a no-op, which would leave the claim set until its lease +// expired even after the sweep finished with the row. +func (s *AuditTransactionStore) ReleaseRecoveryClaim(ctx context.Context, txID string, owner string, message string) error { + return s.claims.releaseClaim(ctx, txID, owner, message) } // GetSchema overrides the base GetSchema to prefix with advisory lock @@ -87,10 +142,9 @@ func NewTransactionStoreWithNotifier(dbs *scommon.RWDB, tableNames sqlcommon.Tab return &TransactionStore{ TransactionStore: commonStore, - readDB: dbs.ReadDB, writeDB: dbs.WriteDB, - tables: tableNames, lockID: createTableLockID("transactions"), + claims: newRecoveryClaimStore(dbs, tableNames), }, nil } @@ -111,16 +165,17 @@ func NewAuditTransactionStore(dbs *scommon.RWDB, tableNames sqlcommon.TableNames TransactionStore: baseStore, writeDB: dbs.WriteDB, lockID: createTableLockID("audittx"), + claims: newRecoveryClaimStore(dbs, tableNames), }, nil } -// ClaimPendingTransactions atomically claims a batch of pending transactions using PostgreSQL's UPDATE...RETURNING. +// claimPending atomically claims a batch of pending transactions using PostgreSQL's UPDATE...RETURNING. // This ensures only one recovery instance can claim a specific transaction. // All state we need lives on the requests table (tx_id PK + stored_at + status // + recovery_claim_* lease columns); the transactions table is no longer // touched. RETURNING tx_id, stored_at directly from the UPDATE removes the // outer join the previous CTE used to recover the timestamp. -func (db *TransactionStore) ClaimPendingTransactions(ctx context.Context, params tokensdriver.RecoveryClaimParams) ([]*tokensdriver.RecoveryClaim, error) { +func (db *recoveryClaimStore) claimPending(ctx context.Context, params tokensdriver.RecoveryClaimParams) ([]*tokensdriver.RecoveryClaim, error) { logger.Debugf("Claiming pending transactions: owner=%s, olderThan=%s, limit=%d, lease=%s", params.Owner, params.OlderThan, params.Limit, params.LeaseDuration) @@ -190,9 +245,9 @@ func (db *TransactionStore) ClaimPendingTransactions(ctx context.Context, params return claimed, nil } -// ReleaseRecoveryClaim releases the recovery claim on a transaction. +// releaseClaim releases the recovery claim on a transaction. // This clears the claim metadata and optionally updates the status message. -func (db *TransactionStore) ReleaseRecoveryClaim(ctx context.Context, txID string, owner string, message string) error { +func (db *recoveryClaimStore) releaseClaim(ctx context.Context, txID string, owner string, message string) error { logger.Debugf("Releasing recovery claim: txID=%s, owner=%s, message=%s", txID, owner, message) // Build the release query using query builder @@ -243,9 +298,9 @@ func (db *TransactionStore) ReleaseRecoveryClaim(ctx context.Context, txID strin return nil } -// CleanupExpiredClaims removes expired recovery claims. +// cleanupExpired removes expired recovery claims. // Returns the number of claims cleaned up. -func (db *TransactionStore) CleanupExpiredClaims(ctx context.Context) (int, error) { +func (db *recoveryClaimStore) cleanupExpired(ctx context.Context) (int, error) { logger.Debug("Cleaning up expired recovery claims") query, args := q.Update(db.tables.Requests).