libsam/pkg/shreder/client.go

package shreder

import (
	"context"
	"fmt"

	"github.com/gagliardetto/solana-go/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

type Client struct {
	conn         *grpc.ClientConn
	client       ShrederServiceClient
	tableLoader  *AddressTables
	subscription map[string]*SubscribeRequestFilterTransactions
}

func NewShrederClient(
	url string,
	rpcClient *rpc.Client,
	subscription map[string]*SubscribeRequestFilterTransactions,
) (*Client, func(), error) {
	if rpcClient == nil {
		return nil, func() {}, fmt.Errorf("rpc client is nil")
	}

	conn, err := grpc.NewClient(url, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, func() {}, err
	}

	s := &Client{
		conn:         conn,
		client:       NewShrederServiceClient(conn),
		subscription: subscription,
		tableLoader:  NewAddressTables(rpcClient),
	}

	return s, func() {
		s.Wait()
	}, nil
}

func (c *Client) Wait() {
	logger.Debug("waiting for shreder client to stop")

	err := c.conn.Close()
	if err != nil {
		logger.Error("failed to close connection: ", "err", err)
	}

	logger.Debug("shreder client stopped")
}

func (c *Client) ReadSync(ctx context.Context, txCh chan<- TxSignalBatch) error {
	stream, err := c.client.SubscribeTransactions(ctx)
	if err != nil {
		return err
	}

	err = stream.Send(&SubscribeTransactionsRequest{
		Transactions: c.subscription,
	})
	if err != nil {
		return err
	}

	for {
		response, err := stream.Recv()
		if err != nil {
			return err
		}

		txBatch := ParseTransaction(response.Transaction, c.tableLoader)
		if len(txBatch) == 0 {
			continue
		}

		// set fixed source for tx signals
		for _, tx := range txBatch {
			tx.Source = "shreder"
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case txCh <- txBatch:
		}
	}
}