Improve Geyser Timelock backup worker with a getProgramAccounts RPC approach (#176)

Author: jeffyanta

ocp/data/blockchain.go

@@ -35,7 +35,7 @@ type BlockchainData interface {
 	GetBlockchainTokenAccountsByOwner(ctx context.Context, account string) ([]ed25519.PublicKey, error)
 	GetBlockchainTransaction(ctx context.Context, sig string, commitment solana.Commitment) (*solana.ConfirmedTransaction, error)
 	GetBlockchainTransactionTokenBalances(ctx context.Context, sig string) (*solana.TransactionTokenBalances, error)
-	GetBlockchainFilteredProgramAccounts(ctx context.Context, program string, offset uint, filterValue []byte) ([]string, uint64, error)
+	GetBlockchainFilteredProgramAccounts(ctx context.Context, program string, offset uint, filterValue []byte) ([]solana.ProgramAccount, uint64, error)
 }
 
 type BlockchainProvider struct {
@@ -320,7 +320,7 @@ func (dp *BlockchainProvider) GetBlockchainTransactionTokenBalances(ctx context.
 	return &res, nil
 }
 
-func (dp *BlockchainProvider) GetBlockchainFilteredProgramAccounts(ctx context.Context, program string, offset uint, filterValue []byte) ([]string, uint64, error) {
+func (dp *BlockchainProvider) GetBlockchainFilteredProgramAccounts(ctx context.Context, program string, offset uint, filterValue []byte) ([]solana.ProgramAccount, uint64, error) {
 	tracer := metrics.TraceMethodCall(ctx, blockchainProviderMetricsName, "GetBlockchainFilteredProgramAccounts")
 	defer tracer.End()
 
@@ -329,10 +329,10 @@ func (dp *BlockchainProvider) GetBlockchainFilteredProgramAccounts(ctx context.C
 		return nil, 0, err
 	}
 
-	addresses, slot, err := dp.sc.GetFilteredProgramAccounts(programId, offset, filterValue)
+	accounts, slot, err := dp.sc.GetFilteredProgramAccounts(programId, offset, filterValue)
 	if err != nil {
 		tracer.OnError(err)
 		return nil, 0, err
 	}
-	return addresses, slot, err
+	return accounts, slot, err
 }

ocp/worker/geyser/backup.go

@@ -5,14 +5,14 @@ import (
 	"sync"
 	"time"
 
+	"github.com/mr-tron/base58"
 	"go.uber.org/zap"
 
-	"github.com/code-payments/ocp-server/database/query"
 	"github.com/code-payments/ocp-server/metrics"
 	"github.com/code-payments/ocp-server/ocp/common"
 	"github.com/code-payments/ocp-server/ocp/data/account"
 	"github.com/code-payments/ocp-server/ocp/data/timelock"
-	timelock_token "github.com/code-payments/ocp-server/solana/timelock/v1"
+	"github.com/code-payments/ocp-server/solana/vm"
 )
 
 // Backup system workers can be found here. This is necessary because we can't rely
@@ -21,7 +21,7 @@ import (
 // the case? Real time updates. Backup workers likely won't be able to guarantee
 // real time (or near real time) updates at scale.
 
-func (p *runtime) backupTimelockStateWorker(runtimeCtx context.Context, state timelock_token.TimelockState, interval time.Duration) error {
+func (p *runtime) backupTimelockStateWorker(runtimeCtx context.Context, interval time.Duration) error {
 	log := p.log.With(zap.String("method", "backupTimelockStateWorker"))
 	log.Debug("worker started")
 
@@ -36,76 +36,74 @@ func (p *runtime) backupTimelockStateWorker(runtimeCtx context.Context, state ti
 		log.Debug("worker stopped")
 	}()
 
-	start := time.Now()
-	cursor := query.EmptyCursor
 	delay := 0 * time.Second // Initially no delay, so we can run right after a deploy
 	for {
 		select {
 		case <-time.After(delay):
-			batchStart := time.Now()
+			start := time.Now()
 
 			func() {
 				provider := runtimeCtx.Value(metrics.ProviderContextKey).(metrics.Provider)
 				trace := provider.StartTrace("geyser_consumer_runtime__backup_timelock_state_worker")
 				defer trace.End()
 				tracedCtx := metrics.NewContext(runtimeCtx, trace)
 
-				timelockRecords, err := p.data.GetAllTimelocksByState(
+				// todo: Also filter on unlock at if the result set gets too large
+				programAccounts, slot, err := p.data.GetBlockchainFilteredProgramAccounts(
 					tracedCtx,
-					state,
-					query.WithDirection(query.Ascending),
-					query.WithCursor(cursor),
-					query.WithLimit(p.conf.backupTimelockWorkerBatchSize.Get(runtimeCtx)),
+					base58.Encode(vm.PROGRAM_ID),
+					0,
+					vm.UnlockStateAccountDiscriminator,
 				)
-				if err == timelock.ErrTimelockNotFound {
-					p.metricStatusLock.Lock()
-					duration := time.Since(start)
-					if p.backupTimelockStateWorkerDuration == nil || *p.backupTimelockStateWorkerDuration < duration {
-						p.backupTimelockStateWorkerDuration = &duration
-					}
-					p.metricStatusLock.Unlock()
-
-					start = time.Now()
-					cursor = query.EmptyCursor
-					return
-				} else if err != nil {
-					log.With(zap.Error(err)).Warn("failed to get timelock records")
+				if err != nil {
+					log.With(zap.Error(err)).Warn("failed to get unlock state program accounts")
 					return
 				}
 
 				reprocessDelay := p.conf.backupTimelockWorkerReprocessDelay.Get(runtimeCtx)
 
-				var wg sync.WaitGroup
-				for _, timelockRecord := range timelockRecords {
-					if lastProcessed, ok := p.backupTimelockProcessedCache.Load(timelockRecord.Address); ok {
+				for _, programAccount := range programAccounts {
+					var unlockState vm.UnlockStateAccount
+					if err := unlockState.Unmarshal(programAccount.Data); err != nil {
+						log.With(zap.Error(err)).Warn("failed to unmarshal unlock state account")
+						continue
+					}
+
+					stateAddress := base58.Encode(unlockState.Address)
+					log := log.With(zap.String("timelock", stateAddress))
+
+					if lastProcessed, ok := p.backupTimelockProcessedCache.Load(stateAddress); ok {
 						if time.Since(lastProcessed.(time.Time)) < reprocessDelay {
 							continue
 						}
 					}
 
-					wg.Add(1)
-
-					go func(timelockRecord *timelock.Record) {
-						defer wg.Done()
-
-						log := log.With(zap.String("timelock", timelockRecord.Address))
+					timelockRecord, err := p.data.GetTimelockByAddress(tracedCtx, stateAddress)
+					if err == timelock.ErrTimelockNotFound {
+						p.backupTimelockProcessedCache.Store(stateAddress, time.Now())
+						continue
+					} else if err != nil {
+						log.With(zap.Error(err)).Warn("failed to get timelock record")
+						continue
+					}
 
-						err := updateTimelockAccountRecord(tracedCtx, p.data, timelockRecord)
-						if err != nil {
-							log.With(zap.Error(err)).Warn("failed to update timelock account")
-							return
-						}
+					if err := updateTimelockAccountRecord(tracedCtx, p.data, timelockRecord, &unlockState, slot); err != nil && err != timelock.ErrStaleTimelockState {
+						log.With(zap.Error(err)).Warn("failed to update timelock account")
+						continue
+					}
 
-						p.backupTimelockProcessedCache.Store(timelockRecord.Address, time.Now())
-					}(timelockRecord)
+					p.backupTimelockProcessedCache.Store(stateAddress, time.Now())
 				}
 
-				wg.Wait()
-
-				cursor = query.ToCursor(timelockRecords[len(timelockRecords)-1].Id)
+				p.metricStatusLock.Lock()
+				duration := time.Since(start)
+				if p.backupTimelockStateWorkerDuration == nil || *p.backupTimelockStateWorkerDuration < duration {
+					p.backupTimelockStateWorkerDuration = &duration
+				}
+				p.metricStatusLock.Unlock()
 			}()
 
-			delay = interval - time.Since(batchStart)
+			delay = interval - time.Since(start)
 		case <-runtimeCtx.Done():
 			return runtimeCtx.Err()
 		}

ocp/worker/geyser/runtime.go

@@ -10,7 +10,6 @@ import (
 	"github.com/code-payments/ocp-server/ocp/integration"
 	"github.com/code-payments/ocp-server/ocp/worker"
 	geyserpb "github.com/code-payments/ocp-server/ocp/worker/geyser/api/gen"
-	timelock_token "github.com/code-payments/ocp-server/solana/timelock/v1"
 
 	ocp_data "github.com/code-payments/ocp-server/ocp/data"
 )
@@ -62,13 +61,7 @@ func New(log *zap.Logger, data ocp_data.Provider, integration integration.Geyser
 func (p *runtime) Start(ctx context.Context, _ time.Duration) error {
 	// Start backup workers to catch missed events
 	go func() {
-		err := p.backupTimelockStateWorker(ctx, timelock_token.StateLocked, p.conf.backupTimelockWorkerInterval.Get(ctx))
-		if err != nil && err != context.Canceled {
-			p.log.With(zap.Error(err)).Warn("timelock backup worker terminated unexpectedly")
-		}
-	}()
-	go func() {
-		err := p.backupTimelockStateWorker(ctx, timelock_token.StateUnknown, p.conf.backupTimelockWorkerInterval.Get(ctx))
+		err := p.backupTimelockStateWorker(ctx, p.conf.backupTimelockWorkerInterval.Get(ctx))
 		if err != nil && err != context.Canceled {
 			p.log.With(zap.Error(err)).Warn("timelock backup worker terminated unexpectedly")
 		}

ocp/worker/geyser/timelock.go

@@ -4,74 +4,30 @@ import (
 	"context"
 	"time"
 
-	"github.com/code-payments/ocp-server/ocp/common"
 	ocp_data "github.com/code-payments/ocp-server/ocp/data"
 	"github.com/code-payments/ocp-server/ocp/data/timelock"
-	"github.com/code-payments/ocp-server/solana"
+	"github.com/code-payments/ocp-server/pointer"
 	timelock_token "github.com/code-payments/ocp-server/solana/timelock/v1"
 	"github.com/code-payments/ocp-server/solana/vm"
 )
 
-func updateTimelockAccountRecord(ctx context.Context, data ocp_data.Provider, timelockRecord *timelock.Record) error {
-	unlockState, slot, err := getTimelockUnlockState(ctx, data, timelockRecord)
-	if err != nil {
-		return err
+func updateTimelockAccountRecord(ctx context.Context, data ocp_data.Provider, timelockRecord *timelock.Record, unlockState *vm.UnlockStateAccount, slot uint64) error {
+	if unlockState == nil {
+		return nil
 	}
 
-	if unlockState != nil {
-		timelockRecord.VaultState = timelock_token.StateWaitingForTimeout
-		if unlockState.IsUnlocked() {
-			timelockRecord.VaultState = timelock_token.StateUnlocked
-		}
+	expectedState := timelock_token.StateWaitingForTimeout
+	if unlockState.IsUnlocked() {
+		expectedState = timelock_token.StateUnlocked
+	}
 
-		unlockAt := uint64(unlockState.UnlockAt)
-		timelockRecord.UnlockAt = &unlockAt
-	} else {
+	if timelockRecord.VaultState == expectedState {
 		return nil
 	}
 
+	timelockRecord.VaultState = expectedState
+	timelockRecord.UnlockAt = pointer.Uint64(uint64(unlockState.UnlockAt))
 	timelockRecord.Block = slot
 	timelockRecord.LastUpdatedAt = time.Now()
 	return data.SaveTimelock(ctx, timelockRecord)
 }
-
-func getTimelockUnlockState(ctx context.Context, data ocp_data.Provider, timelockRecord *timelock.Record) (*vm.UnlockStateAccount, uint64, error) {
-	accountInfoRecord, err := data.GetAccountInfoByTokenAddress(ctx, timelockRecord.VaultAddress)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	vaultOwnerAccount, err := common.NewAccountFromPublicKeyString(timelockRecord.VaultOwner)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	mintAccount, err := common.NewAccountFromPublicKeyString(accountInfoRecord.MintAccount)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	vmConfig, err := common.GetVmConfigForMint(ctx, data, mintAccount)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	timelockAccounts, err := vaultOwnerAccount.GetTimelockAccounts(vmConfig)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	ai, slot, err := data.GetBlockchainAccountInfo(ctx, timelockAccounts.Unlock.PublicKey().ToBase58(), solana.CommitmentFinalized)
-	switch err {
-	case nil:
-		var unlockState vm.UnlockStateAccount
-		if err = unlockState.Unmarshal(ai.Data); err != nil {
-			return nil, 0, err
-		}
-		return &unlockState, slot, nil
-	case solana.ErrNoAccountInfo:
-		return nil, slot, nil
-	default:
-		return nil, 0, err
-	}
-}

solana/client.go

@@ -164,6 +164,11 @@ type TransactionSignature struct {
 	Memo      *string
 }
 
+type ProgramAccount struct {
+	PublicKey ed25519.PublicKey
+	Data      []byte
+}
+
 // Client provides an interaction with the Solana JSON RPC API.
 //
 // Reference: https://docs.solana.com/apps/jsonrpc-api
@@ -177,7 +182,7 @@ type Client interface {
 	GetConfirmationStatus(Signature, Commitment) (bool, error)
 	GetConfirmedBlock(slot uint64) (*Block, error)
 	GetConfirmedBlocksWithLimit(start, limit uint64) ([]uint64, error)
-	GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]string, uint64, error)
+	GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]ProgramAccount, uint64, error)
 	GetLatestBlockhash() (Blockhash, error)
 	GetMinimumBalanceForRentExemption(size uint64) (lamports uint64, err error)
 	GetSignatureStatus(Signature, Commitment) (*SignatureStatus, error)
@@ -1118,7 +1123,7 @@ func (c *client) GetTokenAccountsByOwner(owner, mint ed25519.PublicKey) ([]ed255
 	return keys, nil
 }
 
-func (c *client) GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]string, uint64, error) {
+func (c *client) GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]ProgramAccount, uint64, error) {
 	type memcmpFilter struct {
 		Offset uint   `json:"offset"`
 		Bytes  string `json:"bytes"`
@@ -1152,16 +1157,35 @@ func (c *client) GetFilteredProgramAccounts(program ed25519.PublicKey, offset ui
 			Slot int64 `json:"slot"`
 		} `json:"context"`
 		Value []struct {
-			PubKey string `json:"pubkey"`
+			PubKey  string `json:"pubkey"`
+			Account struct {
+				Data []string `json:"data"`
+			} `json:"account"`
 		} `json:"value"`
 	}
 	if err := c.call(&resp, "getProgramAccounts", base58.Encode(program), config); err != nil {
 		return nil, 0, err
 	}
 
-	var res []string
+	res := make([]ProgramAccount, 0, len(resp.Value))
 	for _, result := range resp.Value {
-		res = append(res, result.PubKey)
+		pubkey, err := base58.Decode(result.PubKey)
+		if err != nil {
+			return nil, 0, errors.Wrap(err, "invalid base58 encoded pubkey")
+		}
+
+		var data []byte
+		if len(result.Account.Data) > 0 {
+			data, err = base64.StdEncoding.DecodeString(result.Account.Data[0])
+			if err != nil {
+				return nil, 0, errors.Wrap(err, "invalid base64 encoded account data")
+			}
+		}
+
+		res = append(res, ProgramAccount{
+			PublicKey: pubkey,
+			Data:      data,
+		})
 	}
 	return res, uint64(resp.Context.Slot), nil
 }

solana/client_with_fallback.go

@@ -140,12 +140,12 @@ func (c *clientWithFallback) GetConfirmedBlocksWithLimit(start, limit uint64) ([
 	)
 }
 
-func (c *clientWithFallback) GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]string, uint64, error) {
+func (c *clientWithFallback) GetFilteredProgramAccounts(program ed25519.PublicKey, offset uint, filterValue []byte) ([]ProgramAccount, uint64, error) {
 	return withFallback2(
-		func() ([]string, uint64, error) {
+		func() ([]ProgramAccount, uint64, error) {
 			return c.primary.GetFilteredProgramAccounts(program, offset, filterValue)
 		},
-		func() ([]string, uint64, error) {
+		func() ([]ProgramAccount, uint64, error) {
 			return c.fallback.GetFilteredProgramAccounts(program, offset, filterValue)
 		},
 	)

solana/client_with_fallback_test.go

@@ -95,7 +95,7 @@ func (m *mockClient) GetConfirmedBlocksWithLimit(uint64, uint64) ([]uint64, erro
 	m.callCount++
 	return nil, nil
 }
-func (m *mockClient) GetFilteredProgramAccounts(ed25519.PublicKey, uint, []byte) ([]string, uint64, error) {
+func (m *mockClient) GetFilteredProgramAccounts(ed25519.PublicKey, uint, []byte) ([]ProgramAccount, uint64, error) {
 	m.callCount++
 	return nil, 0, nil
 }