Improve cleanup of auto-return actions to free up nonces more quickly

Author: jeffyanta

pkg/code/async/account/gift_card.go

@@ -98,14 +98,14 @@ func (p *service) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn
 		log.Trace("gift card is claimed and will be removed from worker queue")
 
 		// Cleanup anything related to gift card auto-return, since it cannot be scheduled
-		err = initiateProcessToCleanupGiftCardAutoReturn(ctx, p.data, giftCardVaultAccount)
+		err = InitiateProcessToCleanupGiftCardAutoReturn(ctx, p.data, giftCardVaultAccount)
 		if err != nil {
 			log.WithError(err).Warn("failure cleaning up auto-return action")
 			return err
 		}
 
 		// Gift card is claimed, so take it out of the worker queue.
-		return markAutoReturnCheckComplete(ctx, p.data, accountInfoRecord)
+		return MarkAutoReturnCheckComplete(ctx, p.data, accountInfoRecord)
 	} else if err != action.ErrActionNotFound {
 		return err
 	}
@@ -129,7 +129,9 @@ func (p *service) maybeInitiateGiftCardAutoReturn(ctx context.Context, accountIn
 		log.WithError(err).Warn("failure initiating process to return gift card balance to issuer")
 		return err
 	}
-	return markAutoReturnCheckComplete(ctx, p.data, accountInfoRecord)
+
+	// Gift card is auto-returned, so take it out of the worker queue
+	return MarkAutoReturnCheckComplete(ctx, p.data, accountInfoRecord)
 }
 
 // Note: This is the first instance of handling a conditional action, and could be
@@ -201,27 +203,30 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data code_data.Pro
 	})
 }
 
-func initiateProcessToCleanupGiftCardAutoReturn(ctx context.Context, data code_data.Provider, giftCardVaultAccount *common.Account) error {
-	autoReturnAction, err := data.GetGiftCardAutoReturnAction(ctx, giftCardVaultAccount.PublicKey().ToBase58())
-	if err != nil {
-		return err
-	}
+// todo: This probably belongs somewhere more common
+func InitiateProcessToCleanupGiftCardAutoReturn(ctx context.Context, data code_data.Provider, giftCardVaultAccount *common.Account) error {
+	return data.ExecuteInTx(ctx, sql.LevelDefault, func(ctx context.Context) error {
+		autoReturnAction, err := data.GetGiftCardAutoReturnAction(ctx, giftCardVaultAccount.PublicKey().ToBase58())
+		if err != nil {
+			return err
+		}
 
-	autoReturnFulfillment, err := data.GetAllFulfillmentsByAction(ctx, autoReturnAction.Intent, autoReturnAction.ActionId)
-	if err != nil {
-		return err
-	}
+		autoReturnFulfillment, err := data.GetAllFulfillmentsByAction(ctx, autoReturnAction.Intent, autoReturnAction.ActionId)
+		if err != nil {
+			return err
+		}
 
-	err = markActionAsRevoked(ctx, data, autoReturnAction)
-	if err != nil {
-		return err
-	}
+		err = markActionAsRevoked(ctx, data, autoReturnAction)
+		if err != nil {
+			return err
+		}
 
-	// The sequencer will handle state transition and any cleanup
-	return markFulfillmentAsActivelyScheduled(ctx, data, autoReturnFulfillment[0])
+		// The sequencer will handle state transition and any cleanup
+		return markFulfillmentAsActivelyScheduled(ctx, data, autoReturnFulfillment[0])
+	})
 }
 
-func markAutoReturnCheckComplete(ctx context.Context, data code_data.Provider, record *account.Record) error {
+func MarkAutoReturnCheckComplete(ctx context.Context, data code_data.Provider, record *account.Record) error {
 	if !record.RequiresAutoReturnCheck {
 		return nil
 	}

pkg/code/async/sequencer/fulfillment_handler.go

@@ -556,7 +556,12 @@ func (h *NoPrivacyWithdrawFulfillmentHandler) OnSuccess(ctx context.Context, ful
 		return err
 	}
 
-	return markTimelockClosed(ctx, h.data, fulfillmentRecord.Source, txnRecord.Slot)
+	err = markTimelockClosed(ctx, h.data, fulfillmentRecord.Source, txnRecord.Slot)
+	if err != nil {
+		return err
+	}
+
+	return maybeCleanupAutoReturnAction(ctx, h.data, fulfillmentRecord.Source)
 }
 
 func (h *NoPrivacyWithdrawFulfillmentHandler) OnFailure(ctx context.Context, fulfillmentRecord *fulfillment.Record, txnRecord *transaction.Record) (recovered bool, err error) {
@@ -752,7 +757,7 @@ func isTokenAccountOnBlockchain(ctx context.Context, data code_data.Provider, ad
 }
 
 // todo: simplify initialization of fulfillment handlers across service and contextual scheduler
-func getFulfillmentHandlers(data code_data.Provider, vmIndexerClient indexerpb.IndexerClient, configProvider ConfigProvider) map[fulfillment.Type]FulfillmentHandler {
+func getFulfillmentHandlers(data code_data.Provider, vmIndexerClient indexerpb.IndexerClient) map[fulfillment.Type]FulfillmentHandler {
 	handlersByType := make(map[fulfillment.Type]FulfillmentHandler)
 	handlersByType[fulfillment.InitializeLockedTimelockAccount] = NewInitializeLockedTimelockAccountFulfillmentHandler(data)
 	handlersByType[fulfillment.NoPrivacyTransferWithAuthority] = NewNoPrivacyTransferWithAuthorityFulfillmentHandler(data, vmIndexerClient)

pkg/code/async/sequencer/scheduler.go

@@ -55,7 +55,7 @@ func NewContextualScheduler(data code_data.Provider, indexerClient indexerpb.Ind
 		log:                     logrus.StandardLogger().WithField("type", "sequencer/scheduler/contextual"),
 		data:                    data,
 		conf:                    configProvider(),
-		handlersByType:          getFulfillmentHandlers(data, indexerClient, configProvider),
+		handlersByType:          getFulfillmentHandlers(data, indexerClient),
 		includeSubsidizerChecks: true,
 	}
 }

pkg/code/async/sequencer/service.go

@@ -40,7 +40,7 @@ func New(data code_data.Provider, scheduler Scheduler, vmIndexerClient indexerpb
 		data:                      data,
 		scheduler:                 scheduler,
 		vmIndexerClient:           vmIndexerClient,
-		fulfillmentHandlersByType: getFulfillmentHandlers(data, vmIndexerClient, configProvider),
+		fulfillmentHandlersByType: getFulfillmentHandlers(data, vmIndexerClient),
 		actionHandlersByType:      getActionHandlers(data),
 		intentHandlersByType:      getIntentHandlers(data),
 	}

pkg/code/async/sequencer/timelock.go

@@ -19,6 +19,10 @@ func markTimelockLocked(ctx context.Context, data code_data.Provider, vault stri
 		return err
 	}
 
+	if record.VaultState == timelock_token_v1.StateLocked {
+		return nil
+	}
+
 	record.VaultState = timelock_token_v1.StateLocked
 	record.Block = slot
 
@@ -35,6 +39,10 @@ func markTimelockClosed(ctx context.Context, data code_data.Provider, vault stri
 		return err
 	}
 
+	if record.VaultState == timelock_token_v1.StateClosed {
+		return nil
+	}
+
 	record.VaultState = timelock_token_v1.StateClosed
 	if record.Block > slot {
 		// Potential conflict with unlock state detection, force a move to close at the next block

pkg/code/async/sequencer/utils.go

@@ -5,7 +5,12 @@ import (
 
 	"github.com/pkg/errors"
 
+	commonpb "github.com/code-payments/code-protobuf-api/generated/go/common/v1"
+
+	async_account "github.com/code-payments/code-server/pkg/code/async/account"
+	"github.com/code-payments/code-server/pkg/code/common"
 	code_data "github.com/code-payments/code-server/pkg/code/data"
+	"github.com/code-payments/code-server/pkg/code/data/action"
 	"github.com/code-payments/code-server/pkg/code/data/fulfillment"
 	"github.com/code-payments/code-server/pkg/code/data/nonce"
 	"github.com/code-payments/code-server/pkg/code/data/transaction"
@@ -291,3 +296,38 @@ func (p *service) markVirtualNonceReleasedDueToSubmittedTransaction(ctx context.
 	nonceRecord.State = nonce.StateReleased
 	return p.data.SaveNonce(ctx, nonceRecord)
 }
+
+func maybeCleanupAutoReturnAction(ctx context.Context, data code_data.Provider, vaultAddress string) error {
+	vaultAccount, err := common.NewAccountFromPublicKeyString(vaultAddress)
+	if err != nil {
+		return err
+	}
+
+	accountInfoRecord, err := data.GetAccountInfoByTokenAddress(ctx, vaultAccount.PublicKey().ToBase58())
+	if err != nil {
+		return err
+	}
+
+	if accountInfoRecord.AccountType != commonpb.AccountType_REMOTE_SEND_GIFT_CARD {
+		return nil
+	}
+
+	_, err = data.GetGiftCardClaimedAction(ctx, vaultAccount.PublicKey().ToBase58())
+	if err == action.ErrActionNotFound {
+		// Gift card isn't claimed, so it must've been auto-returned if it was closed
+		return nil
+	} else if err != nil {
+		return err
+	}
+
+	err = async_account.InitiateProcessToCleanupGiftCardAutoReturn(ctx, data, vaultAccount)
+	if err != nil {
+		return err
+	}
+
+	// It's ok if this fails, the auto-return worker will just process this account
+	// idempotently at a later time
+	async_account.MarkAutoReturnCheckComplete(ctx, data, accountInfoRecord)
+
+	return nil
+}

pkg/code/server/transaction/intent.go

@@ -999,6 +999,10 @@ func (s *transactionServer) VoidGiftCard(ctx context.Context, req *transactionpb
 		return nil, status.Error(codes.Internal, "")
 	}
 
+	// It's ok if this fails, the auto-return worker will just process this account
+	// idempotently at a later time
+	async_account.MarkAutoReturnCheckComplete(ctx, s.data, accountInfoRecord)
+
 	return &transactionpb.VoidGiftCardResponse{
 		Result: transactionpb.VoidGiftCardResponse_OK,
 	}, nil