add option to liteapi

tonkeeper · Dec 5, 2024 · 0a22178 · 0a22178
1 parent c227da0
commit 0a22178
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/liteapi/client.go b/liteapi/client.go
@@ -81,7 +81,8 @@ type Options struct {
 	LiteServers []config.LiteServer
 	Timeout     time.Duration
 	// MaxConnections specifies a number of connections to lite servers for a connections pool.
-	MaxConnections int
+	MaxConnections   int
+	NumPerConnection int
 	// InitCtx is used when opening a new connection to lite servers during the initialization.
 	InitCtx context.Context
 	// ProofPolicy specifies a policy for proof checks.
@@ -114,6 +115,13 @@ func WithMaxConnectionsNumber(maxConns int) Option {
 	}
 }
 
+func WithNumPerConnection(numPerConn int) Option {
+	return func(o *Options) error {
+		o.NumPerConnection = numPerConn
+		return nil
+	}
+}
+
 func WithAsyncConnectionsInit() Option {
 	return func(o *Options) error {
 		o.SyncConnectionsInitialization = false
@@ -269,6 +277,7 @@ func NewClient(options ...Option) (*Client, error) {
 		DetectArchiveNodes:            false,
 		SyncConnectionsInitialization: true,
 		PoolStrategy:                  pool.BestPingStrategy,
+		NumPerConnection:              1,
 	}
 	for _, o := range options {
 		if err := o(opts); err != nil {
@@ -279,7 +288,7 @@ func NewClient(options ...Option) (*Client, error) {
 		return nil, fmt.Errorf("server list empty")
 	}
 	connPool := pool.New(opts.PoolStrategy)
-	initCh := connPool.InitializeConnections(opts.InitCtx, opts.Timeout, opts.MaxConnections, opts.DetectArchiveNodes, opts.LiteServers)
+	initCh := connPool.InitializeConnections(opts.InitCtx, opts.Timeout, opts.MaxConnections, opts.NumPerConnection, opts.DetectArchiveNodes, opts.LiteServers)
 	if opts.SyncConnectionsInitialization {
 		if err := <-initCh; err != nil {
 			return nil, err

diff --git a/liteapi/pool/conn_pool.go b/liteapi/pool/conn_pool.go
@@ -71,13 +71,13 @@ func New(strategy Strategy) *ConnPool {
 	}
 }
 
-func (p *ConnPool) InitializeConnections(ctx context.Context, timeout time.Duration, maxConnections int, detectArchiveNodes bool, servers []config.LiteServer) chan error {
+func (p *ConnPool) InitializeConnections(ctx context.Context, timeout time.Duration, maxConnections int, numPerConnection int, detectArchiveNodes bool, servers []config.LiteServer) chan error {
 	ch := make(chan error, 1)
 	go func() {
 		clientsCh := make(chan clientWrapper, len(servers))
 		for connID, server := range servers {
 			go func(connID int, server config.LiteServer) {
-				cli, _ := connect(ctx, timeout, server)
+				cli, _ := connect(ctx, timeout, server, numPerConnection)
 				// TODO: log error
 				clientsCh <- clientWrapper{
 					connID:     connID,
@@ -119,7 +119,7 @@ func (p *ConnPool) InitializeConnections(ctx context.Context, timeout time.Durat
 	return ch
 }
 
-func connect(ctx context.Context, timeout time.Duration, server config.LiteServer) (*liteclient.Client, error) {
+func connect(ctx context.Context, timeout time.Duration, server config.LiteServer, n int) (*liteclient.Client, error) {
 	serverPubkey, err := base64.StdEncoding.DecodeString(server.Key)
 	if err != nil {
 		return nil, err
@@ -128,7 +128,7 @@ func connect(ctx context.Context, timeout time.Duration, server config.LiteServe
 	if err != nil {
 		return nil, err
 	}
-	cli := liteclient.NewClient(c, liteclient.OptionTimeout(timeout))
+	cli := liteclient.NewClient(c, liteclient.OptionTimeout(timeout), liteclient.OptionConnectionsNum(n))
 	if _, err := cli.LiteServerGetMasterchainInfo(ctx); err != nil {
 		return nil, err
 	}