Create script for reclaiming rent (#586)

Author: rickyrombo

cmd/reclaim_rent/main.go

@@ -0,0 +1,319 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/spf13/cobra"
+
+	"api.audius.co/solana/spl/programs/claimable_tokens"
+	bin "github.com/gagliardetto/binary"
+	"github.com/gagliardetto/solana-go"
+	"github.com/gagliardetto/solana-go/programs/token"
+	"github.com/gagliardetto/solana-go/rpc"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+var reclaimRentCmd = &cobra.Command{
+	Use:   "reclaim_rent <mint>",
+	Short: "Reclaim rent from empty token accounts for a given mint",
+	Args:  cobra.ExactArgs(1),
+	RunE:  reclaimRent,
+}
+
+func main() {
+	err := reclaimRentCmd.Execute()
+	if err != nil {
+		fmt.Println("Error executing command:", err)
+	}
+}
+
+func init() {
+	reclaimRentCmd.Flags().StringP("rpc", "r", "https://api.mainnet-beta.solana.com", "The Solana RPC endpoint to use")
+	reclaimRentCmd.Flags().StringP("database", "c", "postgres://postgres:postgres@localhost:5432/discovery_provider_1?sslmode=disable", "Database connection string")
+	reclaimRentCmd.Flags().StringP("keypair", "k", "~/.config/solana/id.json", "The wallet to use as fee payer for transactions")
+	reclaimRentCmd.Flags().StringP("destination", "d", "", "The recipient of reclaimed rent (defaults to fee payer)")
+	reclaimRentCmd.Flags().StringP("program", "p", claimable_tokens.ProgramID.String(), "The claimable tokens program ID")
+}
+
+func reclaimRent(cmd *cobra.Command, args []string) error {
+	ctx := context.Background()
+
+	rpcEndpoint, err := cmd.Flags().GetString("rpc")
+	if err != nil {
+		return fmt.Errorf("failed to get rpc flag: %w", err)
+	}
+	rpcClient := rpc.New(rpcEndpoint)
+
+	databaseURL, err := cmd.Flags().GetString("database")
+	if err != nil {
+		return fmt.Errorf("failed to get database flag: %w", err)
+	}
+
+	config, err := pgxpool.ParseConfig(databaseURL)
+	if err != nil {
+		return fmt.Errorf("failed to parse database URL: %w", err)
+	}
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	if err != nil {
+		return fmt.Errorf("failed to create database pool: %w", err)
+	}
+	defer pool.Close()
+
+	feePayerFlag, err := cmd.Flags().GetString("keypair")
+	if err != nil {
+		return fmt.Errorf("failed to get keypair flag: %w", err)
+	}
+	keypair, err := solana.PrivateKeyFromSolanaKeygenFile(feePayerFlag)
+	if err != nil {
+		return fmt.Errorf("failed to load keypair: %w", err)
+	}
+
+	destinationFlag, err := cmd.Flags().GetString("destination")
+	if err != nil {
+		return fmt.Errorf("failed to get destination flag: %w", err)
+	}
+	var destination solana.PublicKey
+	if destinationFlag == "" {
+		destination = keypair.PublicKey()
+	} else {
+		destination = solana.MustPublicKeyFromBase58(destinationFlag)
+	}
+
+	programIDFlag, err := cmd.Flags().GetString("program")
+	if err != nil {
+		return fmt.Errorf("failed to get program flag: %w", err)
+	}
+	claimable_tokens.SetProgramID(solana.MustPublicKeyFromBase58(programIDFlag))
+
+	mint := solana.MustPublicKeyFromBase58(args[0])
+
+	fmt.Println("Reclaiming rent for mint:", args[0])
+
+	authority, _, err := claimable_tokens.DeriveAuthority(mint)
+	if err != nil {
+		return fmt.Errorf("failed to derive authority: %w", err)
+	}
+
+	offset := 0
+
+	totalCount, err := getTokenAccountsCountFromDatabase(ctx, pool, mint)
+	if err != nil {
+		return fmt.Errorf("failed to get token accounts count from database: %w", err)
+	}
+
+	limit := 1000
+
+	for {
+		accounts, err := getTokenAccountsFromDatabase(ctx, pool, mint, limit, offset)
+		if err != nil {
+			return fmt.Errorf("failed to get token accounts from database: %w", err)
+		}
+
+		if len(accounts) == 0 {
+			fmt.Println("No more accounts to process.")
+			break
+		}
+		fmt.Printf("Gathered %d accounts from db\n", len(accounts))
+
+		offset += len(accounts)
+
+		filtered, err := filterAccounts(ctx, rpcClient, accounts)
+		if err != nil {
+			return fmt.Errorf("failed to filter accounts: %w", err)
+		}
+
+		batchSize := 15
+		i := 0
+
+		for {
+			batch := make([]DatabaseAccount, 0, batchSize)
+			for j := i; j < i+batchSize && j < len(filtered); j++ {
+				batch = append(batch, filtered[j])
+			}
+			if len(batch) == 0 {
+				break
+			}
+			txSig, err := processBatch(ctx, rpcClient, batch, authority, destination, keypair)
+			if err != nil {
+				return fmt.Errorf("failed to process batch: %w", err)
+			}
+			if txSig != nil {
+				fmt.Printf("Submitted transaction %s to reclaim rent for %d accounts\n", txSig.String(), len(batch))
+			}
+			time.Sleep(time.Second / 500 * 2) // Max 500 req/s (2 req per batch) to avoid rate limiting
+
+			fmt.Printf("Processed %d/%d accounts (%d/%d)\n", i+len(batch), len(filtered), offset/limit+1, (totalCount+999)/limit)
+			i += batchSize
+		}
+	}
+	return nil
+}
+
+func filterAccounts(ctx context.Context, rpcClient *rpc.Client, batch []DatabaseAccount) ([]DatabaseAccount, error) {
+	accounts := make([]solana.PublicKey, 0, len(batch))
+	for _, acct := range batch {
+		accounts = append(accounts, solana.MustPublicKeyFromBase58(acct.Account))
+	}
+
+	res, err := rpcClient.GetMultipleAccountsWithOpts(ctx, accounts, &rpc.GetMultipleAccountsOpts{
+		Encoding: solana.EncodingBase64,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get accounts: %w", err)
+	}
+
+	filtered := make([]DatabaseAccount, 0, len(batch))
+
+	for i, acct := range res.Value {
+		if acct == nil {
+			fmt.Printf("Skipping account %s: account does not exist\n", batch[i].Account)
+			continue
+		}
+		var tokenAccount token.Account
+		err := bin.NewBorshDecoder(acct.Data.GetBinary()).Decode(&tokenAccount)
+		if err != nil {
+			fmt.Printf("Skipping account %s: failed to decode account data (%v)\n", batch[i].Account, err)
+			continue
+		}
+		if tokenAccount.Amount != 0 {
+			fmt.Printf("Skipping account %s: account balance is not zero\n", batch[i].Account)
+			continue
+		}
+		filtered = append(filtered, batch[i])
+	}
+	return filtered, nil
+}
+
+func processBatch(ctx context.Context, rpcClient *rpc.Client, batch []DatabaseAccount, authority solana.PublicKey, destination solana.PublicKey, keypair solana.PrivateKey) (*solana.Signature, error) {
+	instructions := make([]solana.Instruction, 0, len(batch))
+	for _, acct := range batch {
+		closeInstruction := claimable_tokens.NewCloseInstructionBuilder().
+			SetUserBank(solana.MustPublicKeyFromBase58(acct.Account)).
+			SetAuthority(authority).
+			SetDestination(destination).
+			SetEthAddress(common.HexToAddress(acct.EthereumAddress))
+		instructions = append(instructions, closeInstruction.Build())
+	}
+
+	if len(instructions) == 0 {
+		fmt.Println("No valid accounts to process in this batch.")
+		return nil, nil
+	}
+
+	blockhashResult, err := rpcClient.GetLatestBlockhash(ctx, rpc.CommitmentFinalized)
+	if err != nil {
+		return nil, fmt.Errorf("error getting recent blockhash: %v", err)
+	}
+	recentBlockhash := blockhashResult.Value.Blockhash
+
+	tx, err := solana.NewTransaction(
+		instructions,
+		recentBlockhash,
+		solana.TransactionPayer(keypair.PublicKey()),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("error building transaction: %v", err)
+	}
+
+	tx.Sign(func(key solana.PublicKey) *solana.PrivateKey {
+		if key.Equals(keypair.PublicKey()) {
+			return &keypair
+		}
+		return nil
+	})
+
+	txSig := tx.Signatures[0]
+	_, err = rpcClient.SendTransactionWithOpts(ctx, tx, rpc.TransactionOpts{
+		SkipPreflight: true,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("error sending transaction %s: %v", txSig.String(), err)
+	}
+	return &txSig, nil
+}
+
+func getEmptyTokenAccounts(ctx context.Context, client *rpc.Client, mint solana.PublicKey, owner solana.PublicKey, pageKey *string) (rpc.GetProgramAccountsV2Result, error) {
+	mintOffset := uint64(0)
+	ownerOffset := uint64(32)
+	balanceOffset := uint64(64)
+	balance := make([]byte, 8)
+
+	dataSliceOffset := uint64(0)
+	dataSliceLength := uint64(0)
+	limit := uint64(10000)
+
+	return client.GetProgramAccountsV2WithOpts(ctx, solana.TokenProgramID, &rpc.GetProgramAccountsV2Opts{
+		GetProgramAccountsOpts: rpc.GetProgramAccountsOpts{
+			DataSlice: &rpc.DataSlice{
+				Offset: &dataSliceOffset,
+				Length: &dataSliceLength,
+			},
+			Filters: []rpc.RPCFilter{
+				{
+					Memcmp: &rpc.RPCFilterMemcmp{
+						Offset: mintOffset,
+						Bytes:  mint[:],
+					},
+				},
+				{
+					Memcmp: &rpc.RPCFilterMemcmp{
+						Offset: ownerOffset,
+						Bytes:  owner[:],
+					},
+				},
+				{
+					Memcmp: &rpc.RPCFilterMemcmp{
+						Offset: balanceOffset,
+						Bytes:  balance,
+					},
+				},
+				{
+					DataSize: 165, // Standard SPL Token account size
+				},
+			},
+		},
+		PaginationKey: pageKey,
+		Limit:         &limit,
+	})
+}
+
+type DatabaseAccount struct {
+	Account         string
+	EthereumAddress string
+}
+
+func getTokenAccountsFromDatabase(ctx context.Context, pool *pgxpool.Pool, mint solana.PublicKey, limit, offset int) ([]DatabaseAccount, error) {
+	sql := `
+		SELECT bank_account AS account, ethereum_address
+		FROM user_bank_accounts
+		LIMIT $1 OFFSET $2
+	`
+	rows, err := pool.Query(ctx, sql, limit, offset)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query token accounts: %w", err)
+	}
+
+	accounts, err := pgx.CollectRows(rows, pgx.RowToStructByName[DatabaseAccount])
+	if err != nil {
+		return nil, fmt.Errorf("failed to collect token accounts: %w", err)
+	}
+
+	return accounts, nil
+}
+
+func getTokenAccountsCountFromDatabase(ctx context.Context, pool *pgxpool.Pool, mint solana.PublicKey) (int, error) {
+	sql := `
+		SELECT COUNT(*)
+		FROM user_bank_accounts
+	`
+	var count int
+	err := pool.QueryRow(ctx, sql).Scan(&count)
+	if err != nil {
+		return 0, fmt.Errorf("failed to query token accounts count: %w", err)
+	}
+	return count, nil
+}

go.mod

@@ -4,6 +4,7 @@ go 1.25.3
 
 require (
 	connectrpc.com/connect v1.18.1
+	github.com/AlecAivazis/survey/v2 v2.3.7
 	github.com/Doist/unfurlist v0.0.0-20250409100812-515f2735f8e5
 	github.com/OpenAudio/go-openaudio v1.0.9
 	github.com/aquasecurity/esquery v0.2.0
@@ -34,6 +35,7 @@ require (
 	github.com/rpcpool/yellowstone-grpc/examples/golang v0.0.0-20250605231917-29d62ca5d4ae
 	github.com/segmentio/encoding v0.4.1
 	github.com/speps/go-hashids/v2 v2.0.1
+	github.com/spf13/cobra v1.10.1
 	github.com/stretchr/testify v1.11.1
 	github.com/test-go/testify v1.1.4
 	github.com/tidwall/gjson v1.18.0
@@ -49,7 +51,6 @@ require (
 
 require (
 	filippo.io/edwards25519 v1.0.0-rc.1 // indirect
-	github.com/AlecAivazis/survey/v2 v2.3.7 // indirect
 	github.com/DataDog/zstd v1.5.2 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/StackExchange/wmi v1.2.1 // indirect
@@ -204,7 +205,6 @@ require (
 	github.com/shirou/gopsutil/v4 v4.25.1 // indirect
 	github.com/shopspring/decimal v1.4.0 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
-	github.com/spf13/cobra v1.10.1 // indirect
 	github.com/spf13/pflag v1.0.10 // indirect
 	github.com/streamingfast/logging v0.0.0-20230608130331-f22c91403091 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
@@ -243,4 +243,4 @@ require (
 	rsc.io/tmplfunc v0.0.3 // indirect
 )
 
-replace github.com/gagliardetto/solana-go => github.com/rickyrombo/solana-go v1.12.1-0.20250714063439-a36a9577196d
+replace github.com/gagliardetto/solana-go => github.com/rickyrombo/solana-go v0.0.0-20251201234416-e59646f7798f