loading address tables

pkg/shreder/addresstables.go

@@ -0,0 +1,100 @@
+package shreder
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/gagliardetto/solana-go"
+	"github.com/gagliardetto/solana-go/rpc"
+	"github.com/panjf2000/ants/v2"
+)
+
+type AddressTables struct {
+	rpcClient *rpc.Client
+	mux       sync.RWMutex
+	tables    map[solana.PublicKey][]solana.PublicKey
+	loading   map[solana.PublicKey]struct{}
+
+	pool *ants.Pool
+}
+
+func NewAddressTables(rpcClient *rpc.Client) *AddressTables {
+	pool, _ := ants.NewPool(10, ants.WithPreAlloc(true), ants.WithNonblocking(true))
+	return &AddressTables{
+		rpcClient: rpcClient,
+		tables:    make(map[solana.PublicKey][]solana.PublicKey),
+		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[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[tablePubkey] = table
+			total := len(at.tables)
+			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(idx[i]) >= len(addresses) {
+			continue
+		}
+		result = append(result, addresses[i])
+	}
+	return result
+}