universe: fix TCP connection leak by explicitly closing gRPC conn for syncers

In this commit, we attempt to fix a TCP connection leak by explicitly
closing the gRPC connections we create once we're done with the relevant
gRPC client. Otherwise, we'll end up making a new connection for each
new asset to be pushed, which can add up. In the future, we should also
look into the server-side keep alive options.

Author: Roasbeef

universe/auto_syncer.go

@@ -243,6 +243,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/interface.go

@@ -401,6 +401,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
@@ -605,6 +608,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

@@ -462,6 +462,8 @@ func (s *SimpleSyncer) SyncUniverse(ctx context.Context, host ServerAddr,
 			"engine: %w", err)
 	}
 
+	defer diffEngine.Close()
+
 	// With the engine created, we can now sync the local Universe with the
 	// remote instance.
 	return s.executeSync(ctx, diffEngine, syncType, syncConfigs, idsToSync)

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
 }