libsam/pkg/shreder/addresstables.go

package shreder

import (
	"context"
	"fmt"
	"sync"

	"github.com/gagliardetto/solana-go"
	"github.com/gagliardetto/solana-go/rpc"
	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/panjf2000/ants/v2"
)

type AddressTables struct {
	rpcClient *rpc.Client
	mux       sync.RWMutex
	tables    *lru.Cache[solana.PublicKey, []solana.PublicKey]
	loading   map[solana.PublicKey]struct{}

	pool *ants.Pool
}

func NewAddressTables(rpcClient *rpc.Client) *AddressTables {
	pool, _ := ants.NewPool(5, ants.WithPreAlloc(true), ants.WithNonblocking(true))
	cache, _ := lru.New[solana.PublicKey, []solana.PublicKey](10000)
	return &AddressTables{
		rpcClient: rpcClient,
		tables:    cache,
		loading:   make(map[solana.PublicKey]struct{}),
		pool:      pool,
	}
}

func (at *AddressTables) loadAddressTable(tablePubkey solana.PublicKey) ([]solana.PublicKey, error) {
	// decode acc
	acc, err := at.rpcClient.GetAccountInfoWithOpts(context.Background(), tablePubkey, &rpc.GetAccountInfoOpts{
		Encoding: solana.EncodingBase64,
	})
	if err != nil {
		return nil, err
	}
	data := acc.GetBinary()
	if len(data) <= 56 {
		return nil, fmt.Errorf("account data too short")
	}
	offset := 56
	var addresses solana.PublicKeySlice = make([]solana.PublicKey, 0, (len(data)-offset)/32)
	for offset+32 <= len(data) {
		addresses = append(addresses, solana.PublicKeyFromBytes(data[offset:offset+32]))
		offset += 32
	}
	// addresses = append(addresses, solana.PublicKeyFromBytes(data[start:start+32]))
	return addresses, nil

}

func (at *AddressTables) GetAddressTable(tablePubkey solana.PublicKey, idx []uint8) []solana.PublicKey {
	at.mux.RLock()
	addresses, ok := at.tables.Get(tablePubkey)
	if !ok {
		at.mux.RUnlock()
		_ = at.pool.Submit(func() {
			at.mux.RLock()
			_, loading := at.loading[tablePubkey]
			if loading {
				at.mux.RUnlock()
				return
			}
			at.mux.RUnlock()
			at.mux.Lock()
			at.loading[tablePubkey] = struct{}{}
			at.mux.Unlock()

			table, err := at.loadAddressTable(tablePubkey)
			if err != nil {
				logger.Error("loadAddressTable failed", "err", err, "table", tablePubkey)
				at.mux.Lock()
				delete(at.loading, tablePubkey)
				at.mux.Unlock()
				return
			}
			at.mux.Lock()
			at.tables.Add(tablePubkey, table)
			total := at.tables.Len()
			delete(at.loading, tablePubkey)
			at.mux.Unlock()
			logger.Info("loadAddressTable", "table", tablePubkey.String(), "table count:", total)
		})
		return nil
	}
	at.mux.RUnlock()

	var result solana.PublicKeySlice = make([]solana.PublicKey, 0, len(idx))
	for _, i := range idx {
		if int(i) >= len(addresses) {
			continue
		}
		result = append(result, addresses[i])
	}
	return result
}