StickyContextNextNodeBalanced method

liteclient/integration_test.go

@@ -87,13 +87,17 @@ func Test_ConnSticky(t *testing.T) {
 
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
-	ctx = client.StickyContext(ctx)
 
 	err := client.AddConnectionsFromConfigUrl(ctx, "https://tonutils.com/global.config.json")
 	if err != nil {
 		t.Fatal("add connections err", err)
 	}
 
+	ctx, err = client.StickyContextNextNodeBalanced(ctx)
+	if err != nil {
+		t.Fatal("next balanced err", err)
+	}
+
 	doReq := func(expErr error) {
 		var resp tl.Serializable
 		err := client.QueryLiteserver(ctx, GetMasterchainInf{}, &resp)

liteclient/pool.go

@@ -22,6 +22,7 @@ const _StickyCtxUsedNodesKey = "_ton_used_nodes_sticky"
 var (
 	ErrNoActiveConnections = errors.New("no active connections")
 	ErrADNLReqTimeout      = errors.New("adnl request timeout")
+	ErrNoNodesLeft         = errors.New("no more active nodes left")
 )
 
 type OnDisconnectCallback func(addr, key string)
@@ -97,6 +98,7 @@ func (c *ConnectionPool) StickyContext(ctx context.Context) context.Context {
 	return context.WithValue(ctx, _StickyCtxKey, id)
 }
 
+// StickyContextNextNode - select next node in the available list (pseudo random)
 func (c *ConnectionPool) StickyContextNextNode(ctx context.Context) (context.Context, error) {
 	nodeID, _ := ctx.Value(_StickyCtxKey).(uint32)
 	usedNodes, _ := ctx.Value(_StickyCtxUsedNodesKey).([]uint32)
@@ -118,7 +120,47 @@ iter:
 		return context.WithValue(context.WithValue(ctx, _StickyCtxKey, node.id), _StickyCtxUsedNodesKey, usedNodes), nil
 	}
 
-	return ctx, fmt.Errorf("no more active nodes left")
+	return ctx, ErrNoNodesLeft
+}
+
+// StickyContextNextNodeBalanced - select next node based on its weight and availability
+func (c *ConnectionPool) StickyContextNextNodeBalanced(ctx context.Context) (context.Context, error) {
+	nodeID, _ := ctx.Value(_StickyCtxKey).(uint32)
+	usedNodes, _ := ctx.Value(_StickyCtxUsedNodesKey).([]uint32)
+	if nodeID > 0 {
+		usedNodes = append(usedNodes, nodeID)
+	}
+
+	c.nodesMx.RLock()
+	defer c.nodesMx.RUnlock()
+
+	var reqNode *connection
+
+iter:
+	for _, node := range c.activeNodes {
+		for _, usedNode := range usedNodes {
+			if usedNode == node.id {
+				continue iter
+			}
+		}
+
+		if reqNode == nil {
+			reqNode = node
+			continue
+		}
+
+		// select best node on this moment
+		nw, old := atomic.LoadInt64(&node.weight), atomic.LoadInt64(&reqNode.weight)
+		if nw > old || (nw == old && atomic.LoadInt64(&node.lastRespTime) < atomic.LoadInt64(&reqNode.lastRespTime)) {
+			reqNode = node
+		}
+	}
+
+	if reqNode != nil {
+		return context.WithValue(context.WithValue(ctx, _StickyCtxKey, reqNode.id), _StickyCtxUsedNodesKey, usedNodes), nil
+	}
+
+	return ctx, ErrNoNodesLeft
 }
 
 func (c *ConnectionPool) StickyContextWithNodeID(ctx context.Context, nodeId uint32) context.Context {

ton/api.go

@@ -34,6 +34,7 @@ type LiteClient interface {
 	QueryLiteserver(ctx context.Context, payload tl.Serializable, result tl.Serializable) error
 	StickyContext(ctx context.Context) context.Context
 	StickyContextNextNode(ctx context.Context) (context.Context, error)
+	StickyContextNextNodeBalanced(ctx context.Context) (context.Context, error)
 	StickyNodeID(ctx context.Context) uint32
 }
 

ton/dns/integration_test.go

@@ -33,7 +33,7 @@ func TestDNSClient_Resolve(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
 	defer cancel()
 
-	ctx = client.StickyContext(ctx)
+	ctx, _ = client.StickyContextNextNodeBalanced(ctx)
 
 	d, err := cli.Resolve(ctx, "foundation.ton")
 	if err != nil {

ton/retrier.go

@@ -25,7 +25,7 @@ func (w *retryClient) QueryLiteserver(ctx context.Context, payload tl.Serializab
 		tries++
 
 		if err != nil {
-			if !errors.Is(err, liteclient.ErrADNLReqTimeout) && !errors.Is(err, context.DeadlineExceeded){
+			if !errors.Is(err, liteclient.ErrADNLReqTimeout) && !errors.Is(err, context.DeadlineExceeded) {
 				return err
 			}
 
@@ -72,3 +72,7 @@ func (w *retryClient) StickyNodeID(ctx context.Context) uint32 {
 func (w *retryClient) StickyContextNextNode(ctx context.Context) (context.Context, error) {
 	return w.original.StickyContextNextNode(ctx)
 }
+
+func (w *retryClient) StickyContextNextNodeBalanced(ctx context.Context) (context.Context, error) {
+	return w.original.StickyContextNextNodeBalanced(ctx)
+}

ton/timeouter.go

@@ -30,3 +30,7 @@ func (c *timeoutClient) StickyNodeID(ctx context.Context) uint32 {
 func (c *timeoutClient) StickyContextNextNode(ctx context.Context) (context.Context, error) {
 	return c.original.StickyContextNextNode(ctx)
 }
+
+func (c *timeoutClient) StickyContextNextNodeBalanced(ctx context.Context) (context.Context, error) {
+	return c.original.StickyContextNextNodeBalanced(ctx)
+}

ton/waiter.go

@@ -49,3 +49,7 @@ func (w *waiterClient) StickyNodeID(ctx context.Context) uint32 {
 func (w *waiterClient) StickyContextNextNode(ctx context.Context) (context.Context, error) {
 	return w.original.StickyContextNextNode(ctx)
 }
+
+func (w *waiterClient) StickyContextNextNodeBalanced(ctx context.Context) (context.Context, error) {
+	return w.original.StickyContextNextNodeBalanced(ctx)
+}