Merge pull request #719 from Roasbeef/tcp-connection-issue

multi: attempt to fix issue of high number of TCP connections by explicitly closing connections and implementing an idle timer

Author: Roasbeef

proof/courier.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/davecgh/go-spew/spew"
 	"github.com/lightninglabs/lightning-node-connect/hashmailrpc"
 	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/fn"
@@ -66,6 +67,9 @@ type Courier interface {
 	// SetSubscribers sets the set of subscribers that will be notified
 	// of proof courier related events.
 	SetSubscribers(map[uint64]*fn.EventReceiver[fn.Event])
+
+	// Close stops the courier instance.
+	Close() error
 }
 
 // CourierAddr is a fully validated courier address (including protocol specific
@@ -223,6 +227,7 @@ func (h *UniverseRpcCourierAddr) NewCourier(_ context.Context,
 		backoffHandle: backoffHandle,
 		transfer:      cfg.TransferLog,
 		subscribers:   subscribers,
+		rawConn:       conn,
 	}, nil
 }
 
@@ -291,11 +296,16 @@ type ProofMailbox interface {
 	// CleanUp attempts to tear down the mailbox as specified by the passed
 	// sid.
 	CleanUp(ctx context.Context, sid streamID) error
+
+	// Close closes the underlying connection to the hashmail server.
+	Close() error
 }
 
 // HashMailBox is an implementation of the ProofMailbox interface backed by the
 // hashmailrpc.HashMailClient.
 type HashMailBox struct {
+	rawConn *grpc.ClientConn
+
 	client hashmailrpc.HashMailClient
 }
 
@@ -341,7 +351,8 @@ func NewHashMailBox(courierAddr *url.URL) (*HashMailBox,
 	client := hashmailrpc.NewHashMailClient(conn)
 
 	return &HashMailBox{
-		client: client,
+		client:  client,
+		rawConn: conn,
 	}, nil
 }
 
@@ -480,6 +491,11 @@ func (h *HashMailBox) CleanUp(ctx context.Context, sid streamID) error {
 	return err
 }
 
+// Close closes the underlying connection to the hashmail server.
+func (h *HashMailBox) Close() error {
+	return h.rawConn.Close()
+}
+
 // A compile-time assertion to ensure that the HashMailBox meets the
 // ProofMailbox interface.
 var _ ProofMailbox = (*HashMailBox)(nil)
@@ -853,6 +869,8 @@ func (h *HashMailCourier) DeliverProof(ctx context.Context,
 
 	log.Infof("Received ACK from receiver! Cleaning up mailboxes...")
 
+	defer h.Close()
+
 	// Once we receive this ACK, we can clean up our mailbox and also the
 	// receiver's mailbox.
 	if err := h.mailbox.CleanUp(ctx, senderStreamID); err != nil {
@@ -928,6 +946,17 @@ func (h *HashMailCourier) publishSubscriberEvent(event fn.Event) {
 	}
 }
 
+// Close closes the underlying connection to the hashmail server.
+func (h *HashMailCourier) Close() error {
+	if err := h.mailbox.Close(); err != nil {
+		log.Warnf("unable to close mailbox session, "+
+			"recipient=%v: %v", err, spew.Sdump(h.recipient))
+		return err
+	}
+
+	return nil
+}
+
 // BackoffWaitEvent is an event that is sent to a subscriber each time we
 // wait via the Backoff procedure before retrying to deliver a proof to the
 // receiver.
@@ -1030,6 +1059,10 @@ type UniverseRpcCourier struct {
 	// the universe RPC server.
 	client unirpc.UniverseClient
 
+	// rawConn is the raw connection that the courier will use to interact
+	// with the remote gRPC service.
+	rawConn *grpc.ClientConn
+
 	// backoffHandle is a handle to the backoff procedure used in proof
 	// delivery.
 	backoffHandle *BackoffHandler
@@ -1297,6 +1330,11 @@ func (c *UniverseRpcCourier) publishSubscriberEvent(event fn.Event) {
 	}
 }
 
+// Close closes the courier's connection to the remote gRPC service.
+func (c *UniverseRpcCourier) Close() error {
+	return c.rawConn.Close()
+}
+
 // A compile-time assertion to ensure the UniverseRpcCourier meets the
 // proof.Courier interface.
 var _ Courier = (*UniverseRpcCourier)(nil)

server.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	proxy "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"github.com/lightninglabs/lndclient"
@@ -21,6 +22,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/macaroons"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 	"gopkg.in/macaroon-bakery.v2/bakery"
 )
 
@@ -269,6 +271,12 @@ func (s *Server) RunUntilShutdown(mainErrChan <-chan error) error {
 	serverOpts = append(serverOpts, rpcServerOpts...)
 	serverOpts = append(serverOpts, ServerMaxMsgReceiveSize)
 
+	keepAliveParams := keepalive.ServerParameters{
+		MaxConnectionIdle: time.Minute * 2,
+	}
+
+	serverOpts = append(serverOpts, grpc.KeepaliveParams(keepAliveParams))
+
 	grpcServer := grpc.NewServer(serverOpts...)
 	defer grpcServer.Stop()
 

tapfreighter/chain_porter.go

@@ -671,6 +671,8 @@ func (p *ChainPorter) transferReceiverProof(pkg *sendPackage) error {
 				"service handle: %w", err)
 		}
 
+		defer courier.Close()
+
 		// Update courier events subscribers before attempting to
 		// deliver proof.
 		p.subscriberMtx.Lock()

universe/auto_syncer.go

@@ -175,6 +175,11 @@ func (f *FederationEnvoy) Start() error {
 	return nil
 }
 
+// Close frees up any ephemeral resources allocated by the envoy.
+func (f *FederationEnvoy) Close() error {
+	return nil
+}
+
 // Stop stops all active goroutines.
 func (f *FederationEnvoy) Stop() error {
 	f.stopOnce.Do(func() {
@@ -243,6 +248,8 @@ func (f *FederationEnvoy) pushProofToServer(ctx context.Context,
 			"to remote server(%v): %w", addr.HostStr(), err)
 	}
 
+	defer remoteUniverseServer.Close()
+
 	_, err = remoteUniverseServer.UpsertProofLeaf(
 		ctx, uniID, key, leaf,
 	)

universe/base.go

@@ -72,6 +72,11 @@ func NewArchive(cfg ArchiveConfig) *Archive {
 	return a
 }
 
+// Close closes the archive, stopping all goroutines and freeing all resources.
+func (a *Archive) Close() error {
+	return nil
+}
+
 // fetchUniverse returns the base universe instance for the passed identifier.
 // The universe will be loaded in on demand if it has not been seen before.
 func (a *Archive) fetchUniverse(id Identifier) BaseBackend {

universe/interface.go

@@ -446,6 +446,9 @@ type Registrar interface {
 	// UpsertProofLeaf upserts a proof leaf within the target universe tree.
 	UpsertProofLeaf(ctx context.Context, id Identifier, key LeafKey,
 		leaf *Leaf) (*Proof, error)
+
+	// Close is used to shutdown the active registrar instance.
+	Close() error
 }
 
 // Item contains the data fields necessary to insert/update a proof leaf
@@ -650,6 +653,9 @@ type DiffEngine interface {
 	// of diff
 	FetchProofLeaf(ctx context.Context, id Identifier,
 		key LeafKey) ([]*Proof, error)
+
+	// Close is used to shutdown the active diff engine instance.
+	Close() error
 }
 
 // Commitment is an on chain universe commitment. This includes the merkle

universe/syncer.go

@@ -461,6 +461,7 @@ func (s *SimpleSyncer) SyncUniverse(ctx context.Context, host ServerAddr,
 		return nil, fmt.Errorf("unable to create remote diff "+
 			"engine: %w", err)
 	}
+	defer diffEngine.Close()
 
 	// With the engine created, we can now sync the local Universe with the
 	// remote instance.

universe_rpc_diff.go

@@ -15,7 +15,7 @@ import (
 // RpcUniverseDiff is an implementation of the universe.DiffEngine interface
 // that uses an RPC connection to target Universe.
 type RpcUniverseDiff struct {
-	conn unirpc.UniverseClient
+	conn *universeClientConn
 }
 
 // NewRpcUniverseDiff creates a new RpcUniverseDiff instance that dials out to
@@ -210,6 +210,17 @@ func (r *RpcUniverseDiff) FetchProofLeaf(ctx context.Context,
 	return []*universe.Proof{uniProof}, nil
 }
 
+// Close closes the underlying RPC connection to the remote universe server.
+func (r *RpcUniverseDiff) Close() error {
+	if err := r.conn.Close(); err != nil {
+		tapdLog.Warnf("unable to close universe RPC "+
+			"connection: %v", err)
+		return err
+	}
+
+	return nil
+}
+
 // A compile time interface to ensure that RpcUniverseDiff implements the
 // universe.DiffEngine interface.
 var _ universe.DiffEngine = (*RpcUniverseDiff)(nil)

universe_rpc_registrar.go

@@ -18,7 +18,7 @@ import (
 // RpcUniverseRegistrar is an implementation of the universe.Registrar interface
 // that uses an RPC connection to target Universe.
 type RpcUniverseRegistrar struct {
-	conn unirpc.UniverseClient
+	conn *universeClientConn
 }
 
 // NewRpcUniverseRegistrar creates a new RpcUniverseRegistrar instance that
@@ -115,6 +115,17 @@ func (r *RpcUniverseRegistrar) UpsertProofLeaf(ctx context.Context,
 	return unmarshalIssuanceProof(uniKey, proofResp)
 }
 
+// Close closes the underlying RPC connection to the remote Universe server.
+func (r *RpcUniverseRegistrar) Close() error {
+	if err := r.conn.Close(); err != nil {
+		tapdLog.Warnf("unable to close universe RPC "+
+			"connection: %v", err)
+		return err
+	}
+
+	return nil
+}
+
 // A compile time interface to ensure that RpcUniverseRegistrar implements the
 // universe.Registrar interface.
 var _ universe.Registrar = (*RpcUniverseRegistrar)(nil)
@@ -151,10 +162,18 @@ func CheckFederationServer(localRuntimeID int64, connectTimeout time.Duration,
 	return nil
 }
 
+// universeClientConn is a wrapper around a gRPC client connection that also
+// includes the raw connection. This allows us to properly manage the lifecycle
+// of the connection.
+type universeClientConn struct {
+	*grpc.ClientConn
+	unirpc.UniverseClient
+}
+
 // ConnectUniverse connects to a remote Universe server using the provided
 // server address.
 func ConnectUniverse(
-	serverAddr universe.ServerAddr) (unirpc.UniverseClient, error) {
+	serverAddr universe.ServerAddr) (*universeClientConn, error) {
 
 	// TODO(roasbeef): all info is authenticated, but also want to allow
 	// brontide connect as well, can avoid TLS certs
@@ -179,5 +198,8 @@ func ConnectUniverse(
 			"server: %v", err)
 	}
 
-	return unirpc.NewUniverseClient(rawConn), nil
+	return &universeClientConn{
+		ClientConn:     rawConn,
+		UniverseClient: unirpc.NewUniverseClient(rawConn),
+	}, nil
 }