From 0b717e5f391f0e78f5c937b654ee27a2a566a708 Mon Sep 17 00:00:00 2001
From: Claudio Netto <nettinhorama@gmail.com>
Date: Tue, 12 May 2020 20:08:36 -0300
Subject: [PATCH] fix: open a new connection for each probe call

---
 cmd/livenessprobe/livenessprobe_test.go | 26 +++++++------------------
 cmd/livenessprobe/main.go               | 24 +++++++++++++++++++----
 2 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/cmd/livenessprobe/livenessprobe_test.go b/cmd/livenessprobe/livenessprobe_test.go
index 38435338..989b6f5a 100644
--- a/cmd/livenessprobe/livenessprobe_test.go
+++ b/cmd/livenessprobe/livenessprobe_test.go
@@ -17,15 +17,14 @@ limitations under the License.
 package main
 
 import (
+	"flag"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	csi "github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/golang/mock/gomock"
-	connlib "github.com/kubernetes-csi/csi-lib-utils/connection"
 	"github.com/kubernetes-csi/csi-test/driver"
-	"google.golang.org/grpc"
 )
 
 const (
@@ -37,9 +36,7 @@ func createMockServer(t *testing.T) (
 	*driver.MockCSIDriver,
 	*driver.MockIdentityServer,
 	*driver.MockControllerServer,
-	*driver.MockNodeServer,
-	*grpc.ClientConn,
-	error) {
+	*driver.MockNodeServer) {
 	// Start the mock server
 	mockController := gomock.NewController(t)
 	identityServer := driver.NewMockIdentityServer(mockController)
@@ -52,24 +49,16 @@ func createMockServer(t *testing.T) (
 	})
 	drv.Start()
 
-	// Create a client connection to it
-	addr := drv.Address()
-	csiConn, err := connlib.Connect(addr)
-	if err != nil {
-		return nil, nil, nil, nil, nil, nil, err
-	}
-
-	return mockController, drv, identityServer, controllerServer, nodeServer, csiConn, nil
+	return mockController, drv, identityServer, controllerServer, nodeServer
 }
 
 func TestProbe(t *testing.T) {
-	mockController, driver, idServer, _, _, csiConn, err := createMockServer(t)
-	if err != nil {
-		t.Fatal(err)
-	}
+	mockController, driver, idServer, _, _ := createMockServer(t)
 	defer mockController.Finish()
 	defer driver.Stop()
-	defer csiConn.Close()
+
+	flag.Set("csi-address", driver.Address())
+	flag.Parse()
 
 	var injectedErr error
 
@@ -78,7 +67,6 @@ func TestProbe(t *testing.T) {
 	idServer.EXPECT().Probe(gomock.Any(), inProbe).Return(outProbe, injectedErr).Times(1)
 
 	hp := &healthProbe{
-		conn:       csiConn,
 		driverName: driverName,
 	}
 
diff --git a/cmd/livenessprobe/main.go b/cmd/livenessprobe/main.go
index e0e310b5..5e25844b 100644
--- a/cmd/livenessprobe/main.go
+++ b/cmd/livenessprobe/main.go
@@ -39,7 +39,6 @@ var (
 )
 
 type healthProbe struct {
-	conn       *grpc.ClientConn
 	driverName string
 }
 
@@ -47,8 +46,17 @@ func (h *healthProbe) checkProbe(w http.ResponseWriter, req *http.Request) {
 	ctx, cancel := context.WithTimeout(req.Context(), *probeTimeout)
 	defer cancel()
 
+	conn, err := acquireConnection()
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(err.Error()))
+		klog.Errorf("failed to establish connection to CSI driver: %v", err)
+		return
+	}
+	defer conn.Close()
+
 	klog.V(5).Infof("Sending probe request to CSI driver %q", h.driverName)
-	ready, err := rpc.Probe(ctx, h.conn)
+	ready, err := rpc.Probe(ctx, conn)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		w.Write([]byte(err.Error()))
@@ -68,12 +76,20 @@ func (h *healthProbe) checkProbe(w http.ResponseWriter, req *http.Request) {
 	klog.V(5).Infof("Health check succeeded")
 }
 
+// acquireConnection wraps the connlib.Connect func.
+//
+// NOTE: always open a new connection to avoid the issue reported at
+// https://github.com/kubernetes-csi/livenessprobe/issues/68.
+func acquireConnection() (*grpc.ClientConn, error) {
+	return connlib.Connect(*csiAddress)
+}
+
 func main() {
 	klog.InitFlags(nil)
 	flag.Set("logtostderr", "true")
 	flag.Parse()
 
-	csiConn, err := connlib.Connect(*csiAddress)
+	csiConn, err := acquireConnection()
 	if err != nil {
 		// connlib should retry forever so a returned error should mean
 		// the grpc client is misconfigured rather than an error on the network
@@ -82,13 +98,13 @@ func main() {
 
 	klog.Infof("calling CSI driver to discover driver name")
 	csiDriverName, err := rpc.GetDriverName(context.Background(), csiConn)
+	csiConn.Close()
 	if err != nil {
 		klog.Fatalf("failed to get CSI driver name: %v", err)
 	}
 	klog.Infof("CSI driver name: %q", csiDriverName)
 
 	hp := &healthProbe{
-		conn:       csiConn,
 		driverName: csiDriverName,
 	}