close wallet connections methods

connectionprovider.go

@@ -45,12 +45,16 @@ type PuddleConnectionProvider struct {
 	credentials     credentials.TransportCredentials
 }
 
+func (c *PuddleConnectionProvider) getConnectionKey(endpoint *Endpoint) string {
+	return fmt.Sprintf("%s:%s", endpoint.String(), c.name)
+}
+
 // CloseConnections closes connections to the given endpoints and specific connection provider.
 func (c *PuddleConnectionProvider) CloseConnections(endpoints []*Endpoint) {
 	connectionPoolsMu.Lock()
 	defer connectionPoolsMu.Unlock()
 	for i := range endpoints {
-		key := fmt.Sprintf("%s:%s", endpoints[i].String(), c.name)
+		key := c.getConnectionKey(endpoints[i])
 		if pool, exists := connectionPools[key]; exists {
 			pool.Close()
 			delete(connectionPools, key)
@@ -60,7 +64,7 @@ func (c *PuddleConnectionProvider) CloseConnections(endpoints []*Endpoint) {
 
 // Connection returns a connection and release function.
 func (c *PuddleConnectionProvider) Connection(ctx context.Context, endpoint *Endpoint) (*grpc.ClientConn, func(), error) {
-	pool := c.obtainOrCreatePool(fmt.Sprintf("%s:%s", endpoint.String(), c.name))
+	pool := c.obtainOrCreatePool(fmt.Sprintf("%s:%s", endpoint.String(), c.name), c.getConnectionKey(endpoint))
 
 	res, err := pool.Acquire(ctx)
 	if err != nil {
@@ -70,9 +74,9 @@ func (c *PuddleConnectionProvider) Connection(ctx context.Context, endpoint *End
 	return res.Value(), res.Release, nil
 }
 
-func (c *PuddleConnectionProvider) obtainOrCreatePool(address string) *puddle.Pool[*grpc.ClientConn] {
+func (c *PuddleConnectionProvider) obtainOrCreatePool(address, connectionKey string) *puddle.Pool[*grpc.ClientConn] {
 	connectionPoolsMu.Lock()
-	pool, exists := connectionPools[address]
+	pool, exists := connectionPools[connectionKey]
 	connectionPoolsMu.Unlock()
 	if !exists {
 		constructor := func(ctx context.Context) (*grpc.ClientConn, error) {
@@ -107,7 +111,7 @@ func (c *PuddleConnectionProvider) obtainOrCreatePool(address string) *puddle.Po
 			MaxSize:     c.poolConnections,
 		})
 		connectionPoolsMu.Lock()
-		connectionPools[address] = pool
+		connectionPools[connectionKey] = pool
 		connectionPoolsMu.Unlock()
 	}