connection: Disconnect idle gRPC connections on inactivity

This patch adds support to close the lingering gRPC connections
that have not been used for a certain amount of time, default to 5 mins.

An interceptor is used to track the last access time in a reliable
manner.

If a expired connection is requested, it will first be re-created and
then returned to the caller.

Since gRPC connections are not simple disconnects but a complete
tear down, re-creating the connections is necessary.

Signed-off-by: Niraj Yadav <niryadav@redhat.com>

Author: black-dragon74

internal/connection/connection.go

@@ -19,6 +19,8 @@ package connection
 import (
 	"context"
 	"crypto/tls"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/csi-addons/kubernetes-csi-addons/internal/kubernetes/token"
@@ -29,23 +31,67 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 )
 
+const (
+	// The duration after which a gRPC connection is closed due to inactivity
+	idleTimeout = time.Minute * 5
+)
+
 // Connection struct consists of to NodeID, DriverName, Capabilities for the controller
 // to pick sidecar connection and grpc Client to connect to the sidecar.
 type Connection struct {
+	sync.Mutex
+
 	Client       *grpc.ClientConn
 	Capabilities []*identity.Capability
 	Namespace    string
 	Name         string
 	NodeID       string
 	DriverName   string
 	Timeout      time.Duration
+
+	// Holds the internal state of the connection
+	connected  bool
+	enableAuth bool
+	endpoint   string
+	podName    string
+
+	// Used to cancel any existing timers in case of re-connects
+	cancelIdle context.CancelFunc
+	// Holds the last access time of the gRPC Connection
+	// It is tracked and updated by the `accessTimeInterceptor`
+	lastAccessTime atomic.Int64
 }
 
-// NewConnection establishes connection with sidecar, fetches capability and returns Connection object
-// filled with required information.
-func NewConnection(ctx context.Context, endpoint, nodeID, driverName, namespace, podName string, enableAuth bool) (*Connection, error) {
-	var opts []grpc.DialOption
-	if enableAuth {
+// accessTimeInterceptor is an unary interceptor which updates the lastAccessTime
+// on `Connection` struct to `time.Now()` on each RPC call
+func accessTimeInterceptor(conn *Connection) grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		conn.lastAccessTime.Store(time.Now().UnixNano())
+
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+}
+
+// Connect creates a new grpc.ClientConn object and sets it as the
+// client property on Connection struct. If a connection is already
+// connected, it is reused as is. It also spawns a go routine to tear
+// down the connection if it has been idling for a specified threshold.
+//
+// In cases where a new connection is created from the scratch Connect
+// also calls fetchCapabilities on the connection object.
+func (c *Connection) Connect() error {
+	c.Lock()
+	defer c.Unlock()
+
+	// Return early if already connected, the connection will be reused
+	if c.connected {
+		return nil
+	}
+
+	opts := []grpc.DialOption{
+		grpc.WithUnaryInterceptor(accessTimeInterceptor(c)),
+	}
+	if c.enableAuth {
 		opts = append(opts, token.WithServiceAccountToken())
 		tlsConfig := &tls.Config{
 			// Certs are only used to initiate HTTPS connections; authorization is handled by SA tokens
@@ -56,34 +102,110 @@ func NewConnection(ctx context.Context, endpoint, nodeID, driverName, namespace,
 	} else {
 		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	}
-	cc, err := grpc.NewClient(endpoint, opts...)
+	cc, err := grpc.NewClient(c.endpoint, opts...)
 	if err != nil {
-		return nil, err
+		return err
+	}
+	c.Client = cc
+	c.connected = true
+
+	// Fetch the caps
+	// Downside is we need to do this on every re-connect
+	if err := c.fetchCapabilities(context.Background()); err != nil {
+		if e := c.Client.Close(); e == nil {
+			c.connected = false
+		}
+
+		return err
+	}
+
+	// Start a goroutine to close the connection after idle timeout
+	// But first, expire any existing timers
+	if c.cancelIdle != nil {
+		c.cancelIdle()
 	}
+	idleCtx, cFunc := context.WithCancel(context.Background())
+	c.cancelIdle = cFunc
+	go c.startIdleTimer(idleCtx)
+
+	return nil
+}
+
+// startIdleTimer starts a ticker with an interval of 30 seconds
+// At each tick, it checks if the connection has been idle for
+// more than `idleTimeout`. If so, the connection is closed.
+func (c *Connection) startIdleTimer(ctx context.Context) {
+	ticker := time.NewTicker(time.Second * 30)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			c.Lock()
+
+			lastAccess := time.Unix(0, c.lastAccessTime.Load())
+			isIdle := time.Since(lastAccess) > idleTimeout
+			if isIdle && c.connected {
+				// It's okay if there's an error in tearing down the connection
+				// It will be reused by subsequent requests, see Connect() for details
+				if err := c.Client.Close(); err == nil {
+					c.connected = false
+				}
+			}
+			c.Unlock()
+
+		case <-ctx.Done():
+			// Timer was cancelled
+			return
+		}
+	}
+}
+
+// NewConnection establishes connection with sidecar, fetches capability and returns Connection object
+// filled with required information.
+func NewConnection(ctx context.Context, endpoint, nodeID, driverName, namespace, podName string, enableAuth bool) (*Connection, error) {
 
 	conn := &Connection{
-		Client:     cc,
 		Namespace:  namespace,
 		Name:       podName,
 		NodeID:     nodeID,
 		DriverName: driverName,
 		Timeout:    time.Minute,
+
+		endpoint:   endpoint,
+		podName:    podName,
+		enableAuth: enableAuth,
 	}
 
-	err = conn.fetchCapabilities(ctx)
-	if err != nil {
+	if err := conn.Connect(); err != nil {
 		return nil, err
 	}
 
 	return conn, nil
 }
 
+// Close tears down the gRPC connection and terminates the goroutines
+// monitoring the idle timeout by calling cancelIdle()
 func (c *Connection) Close() error {
-	if c.Client == nil {
-		return nil
+	c.Lock()
+	defer c.Unlock()
+
+	if c.cancelIdle != nil {
+		// Cancel the context as well
+		c.cancelIdle()
+		c.cancelIdle = nil
 	}
 
-	return c.Client.Close()
+	if c.Client != nil {
+		err := c.Client.Close()
+		if err == nil {
+			c.connected = false
+		}
+
+		return err
+	}
+
+	return nil
 }
 
 // fetchCapabilities fetches the capability of the connected CSI driver.

internal/connection/connection_pool.go

@@ -28,6 +28,8 @@ import (
 	"github.com/csi-addons/kubernetes-csi-addons/internal/util"
 )
 
+const failedToReconnectFmtStr = "failed to reconnect an inactive connection due to error: %w"
+
 //+kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch
 
 // ConnectionPool consists of map of Connection objects and
@@ -82,34 +84,45 @@ func (cp *ConnectionPool) Delete(key string) {
 
 // getByDriverName returns map of connections filtered by driverName. This function
 // must be called with read lock held.
-func (cp *ConnectionPool) getByDriverName(driverName string) map[string]*Connection {
+func (cp *ConnectionPool) getByDriverName(driverName string) (map[string]*Connection, error) {
 	newPool := make(map[string]*Connection)
 	for k, v := range cp.pool {
 		if v.DriverName != driverName {
 			continue
 		}
+		if err := v.Connect(); err != nil {
+			return nil, fmt.Errorf(failedToReconnectFmtStr, err)
+		}
 		newPool[k] = v
 	}
 
-	return newPool
+	return newPool, nil
 }
 
 // GetByNodeID returns map of connections, filtered with given driverName and optional nodeID.
-func (cp *ConnectionPool) GetByNodeID(driverName, nodeID string) map[string]*Connection {
+func (cp *ConnectionPool) GetByNodeID(driverName, nodeID string) (map[string]*Connection, error) {
 	cp.rwlock.RLock()
 	defer cp.rwlock.RUnlock()
 
-	pool := cp.getByDriverName(driverName)
+	pool, err := cp.getByDriverName(driverName)
+	if err != nil {
+		return nil, err
+	}
 	result := make(map[string]*Connection)
 	for k, v := range pool {
 		// since nodeID is options,check only if it is not empty
 		if nodeID != "" && v.NodeID != nodeID {
 			continue
 		}
+		// We may skip this one as the validation is done in `getByDriverName` as well
+		// But since this is just checking a boolean, leaving it as-is is fine too.
+		if err = v.Connect(); err != nil {
+			return nil, fmt.Errorf(failedToReconnectFmtStr, err)
+		}
 		result[k] = v
 	}
 
-	return result
+	return result, nil
 }
 
 // getNamespaceByDriverName loops through the connections in the pool and
@@ -166,5 +179,9 @@ func (cp *ConnectionPool) GetLeaderByDriver(ctx context.Context, reconciler clie
 		return nil, fmt.Errorf("no connection with key %q found for driver %q: %w", key, driverName, err)
 	}
 
+	if err = conn.Connect(); err != nil {
+		return nil, fmt.Errorf(failedToReconnectFmtStr, err)
+	}
+
 	return conn, nil
 }