livenessprobe should use CSI connection lib

Author: andrewsykim

cmd/livenessprobe_test.go

@@ -53,7 +53,7 @@ func createMockServer(t *testing.T) (
 
 	// Create a client connection to it
 	addr := drv.Address()
-	csiConn, err := connection.NewConnection(addr, 10)
+	csiConn, err := connection.NewConnection(addr)
 	if err != nil {
 		return nil, nil, nil, nil, nil, nil, err
 	}

cmd/main.go

@@ -36,7 +36,7 @@ const (
 // Command line flags
 var (
 	// kubeconfig        = flag.String("kubeconfig", "", "Absolute path to the kubeconfig file. Required only when running out of cluster.")
-	connectionTimeout = flag.Duration("connection-timeout", 30*time.Second, "Timeout for waiting for CSI driver socket in seconds.")
+	connectionTimeout = flag.Duration("connection-timeout", 0, "The --connection-timeout flag is deprecated")
 	csiAddress        = flag.String("csi-address", "/run/csi/socket", "Address of the CSI driver socket.")
 	healthzPort       = flag.String("health-port", "9808", "TCP ports for listening healthz requests")
 )
@@ -56,10 +56,8 @@ func runProbe(ctx context.Context, csiConn connection.CSIConnection) error {
 }
 
 func getCSIConnection() (connection.CSIConnection, error) {
-	// Connect to CSI.
 	klog.Infof("Attempting to open a gRPC connection with: %s", *csiAddress)
-	csiConn, err := connection.NewConnection(*csiAddress, *connectionTimeout)
-	return csiConn, err
+	return connection.NewConnection(*csiAddress)
 }
 
 func checkHealth(w http.ResponseWriter, req *http.Request) {
@@ -72,7 +70,7 @@ func checkHealth(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 	defer csiConn.Close()
-	ctx, cancel := context.WithTimeout(context.Background(), *connectionTimeout)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	if err := runProbe(ctx, csiConn); err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
@@ -90,6 +88,10 @@ func main() {
 	flag.Set("logtostderr", "true")
 	flag.Parse()
 
+	if *connectionTimeout != 0 {
+		klog.Warning("--connection-timeout is deprecated and will have no effect")
+	}
+
 	addr := net.JoinHostPort("0.0.0.0", *healthzPort)
 	http.HandleFunc("/healthz", checkHealth)
 	klog.Infof("Serving requests to /healthz on: %s", addr)

pkg/connection/connection.go

@@ -19,14 +19,11 @@ package connection
 import (
 	"context"
 	"fmt"
-	"net"
-	"strings"
-	"time"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
+	"github.com/kubernetes-csi/csi-lib-utils/connection"
 	"github.com/kubernetes-csi/csi-lib-utils/protosanitizer"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/connectivity"
 	"k8s.io/klog"
 )
 
@@ -53,9 +50,8 @@ var (
 )
 
 // NewConnection establishes new connection to CSI driver
-func NewConnection(
-	address string, timeout time.Duration) (CSIConnection, error) {
-	conn, err := connect(address, timeout)
+func NewConnection(address string) (CSIConnection, error) {
+	conn, err := connection.Connect(address)
 	if err != nil {
 		return nil, err
 	}
@@ -64,38 +60,6 @@ func NewConnection(
 	}, nil
 }
 
-func connect(address string, timeout time.Duration) (*grpc.ClientConn, error) {
-	klog.Infof("Connecting to %s", address)
-	dialOptions := []grpc.DialOption{
-		grpc.WithInsecure(),
-		grpc.WithBackoffMaxDelay(time.Second),
-		grpc.WithUnaryInterceptor(logGRPC),
-	}
-	if strings.HasPrefix(address, "/") {
-		dialOptions = append(dialOptions, grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
-			return net.DialTimeout("unix", addr, timeout)
-		}))
-	}
-	conn, err := grpc.Dial(address, dialOptions...)
-
-	if err != nil {
-		return nil, err
-	}
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	defer cancel()
-	for {
-		if !conn.WaitForStateChange(ctx, conn.GetState()) {
-			klog.Infof("Connection timed out")
-			return conn, nil // return nil, subsequent GetPluginInfo will show the real connection error
-		}
-		if conn.GetState() == connectivity.Ready {
-			klog.Infof("Connected")
-			return conn, nil
-		}
-		klog.Infof("Still trying, connection is %s", conn.GetState())
-	}
-}
-
 func (c *csiConnection) GetDriverName(ctx context.Context) (string, error) {
 	client := csi.NewIdentityClient(c.conn)