Add client.GetPoolStatus()

tonkeeper · Nov 29, 2024 · 4bfca28 · 4bfca28
1 parent 40dcba5
commit 4bfca28
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 8 deletions.
diff --git a/liteapi/client.go b/liteapi/client.go
@@ -1089,3 +1089,7 @@ func (c *Client) GetNetworkGlobalID(ctx context.Context) (int32, error) {
 	c.networkGlobalID = &block.GlobalId
 	return block.GlobalId, nil
 }
+
+func (c *Client) GetPoolStatus() pool.Status {
+	return c.pool.Status()
+}
diff --git a/liteapi/pool/conn_pool.go b/liteapi/pool/conn_pool.go
@@ -58,6 +58,7 @@ type conn interface {
 	Run(ctx context.Context, detectArchive bool)
 	IsArchiveNode() bool
 	AverageRoundTrip() time.Duration
+	Status() ConnStatus
 }
 
 // New returns a new instance of a connections pool.
@@ -79,8 +80,9 @@ func (p *ConnPool) InitializeConnections(ctx context.Context, timeout time.Durat
 				cli, _ := connect(ctx, timeout, server)
 				// TODO: log error
 				clientsCh <- clientWrapper{
-					connID: connID,
-					cli:    cli,
+					connID:     connID,
+					cli:        cli,
+					serverHost: server.Host,
 				}
 			}(connID, server)
 		}
@@ -99,7 +101,7 @@ func (p *ConnPool) InitializeConnections(ctx context.Context, timeout time.Durat
 					continue
 				}
 				if p.ConnectionsNumber() < maxConnections {
-					c := p.addConnection(wrapper.connID, wrapper.cli)
+					c := p.addConnection(wrapper.connID, wrapper.cli, wrapper.serverHost)
 					go c.Run(context.TODO(), detectArchiveNodes)
 				}
 				if p.ConnectionsNumber() == maxConnections {
@@ -134,15 +136,17 @@ func connect(ctx context.Context, timeout time.Duration, server config.LiteServe
 }
 
 type clientWrapper struct {
-	connID int
-	cli    *liteclient.Client
+	connID     int
+	cli        *liteclient.Client
+	serverHost string
 }
 
-func (p *ConnPool) addConnection(connID int, cli *liteclient.Client) *connection {
+func (p *ConnPool) addConnection(connID int, cli *liteclient.Client, serverHost string) *connection {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	c := &connection{
 		id:                  connID,
+		serverHost:          serverHost,
 		client:              cli,
 		masterHeadUpdatedCh: p.masterHeadUpdatedCh,
 	}
@@ -369,3 +373,20 @@ func (p *ConnPool) WaitMasterchainSeqno(ctx context.Context, seqno uint32, timeo
 		}
 	}
 }
+
+type Status struct {
+	Connections []ConnStatus
+}
+
+func (p *ConnPool) Status() Status {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	connStatuses := make([]ConnStatus, 0, len(p.conns))
+	for _, c := range p.conns {
+		connStatuses = append(connStatuses, c.Status())
+	}
+	return Status{
+		Connections: connStatuses,
+	}
+}
diff --git a/liteapi/pool/conn_pool_test.go b/liteapi/pool/conn_pool_test.go
@@ -47,6 +47,10 @@ func (m *mockConn) Client() *liteclient.Client {
 	panic("implement me")
 }
 
+func (m *mockConn) Status() ConnStatus {
+	panic("implement me")
+}
+
 func (m *mockConn) Run(ctx context.Context, detectArchiveNodes bool) {
 }
 

diff --git a/liteapi/pool/connection.go b/liteapi/pool/connection.go
@@ -10,8 +10,9 @@ import (
 )
 
 type connection struct {
-	id     int
-	client *liteclient.Client
+	id         int
+	serverHost string
+	client     *liteclient.Client
 
 	// masterHeadUpdatedCh is used to send a notification when a known master head is changed.
 	masterHeadUpdatedCh chan masterHeadUpdated
@@ -154,3 +155,17 @@ func (c *connection) setArchive(archive bool) {
 func (c *connection) AverageRoundTrip() time.Duration {
 	return c.client.AverageRoundTrip()
 }
+
+type ConnStatus struct {
+	ServerHost string
+	Connected  bool
+	Archive    bool
+}
+
+func (c *connection) Status() ConnStatus {
+	return ConnStatus{
+		ServerHost: c.serverHost,
+		Connected:  c.IsOK(),
+		Archive:    c.IsArchiveNode(),
+	}
+}