livenessprobe should only connect on startup & only probe on each health check

Author: andrewsykim

cmd/livenessprobe_test.go

@@ -17,14 +17,15 @@ limitations under the License.
 package main
 
 import (
-	"context"
+	"net/http"
+	"net/http/httptest"
 	"testing"
-	"time"
 
 	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"
-	"github.com/kubernetes-csi/livenessprobe/pkg/connection"
+	"google.golang.org/grpc"
 )
 
 const (
@@ -37,7 +38,7 @@ func createMockServer(t *testing.T) (
 	*driver.MockIdentityServer,
 	*driver.MockControllerServer,
 	*driver.MockNodeServer,
-	connection.CSIConnection,
+	*grpc.ClientConn,
 	error) {
 	// Start the mock server
 	mockController := gomock.NewController(t)
@@ -53,7 +54,7 @@ func createMockServer(t *testing.T) (
 
 	// Create a client connection to it
 	addr := drv.Address()
-	csiConn, err := connection.NewConnection(addr)
+	csiConn, err := connlib.Connect(addr)
 	if err != nil {
 		return nil, nil, nil, nil, nil, nil, err
 	}
@@ -72,20 +73,28 @@ func TestProbe(t *testing.T) {
 
 	var injectedErr error
 
-	// Setting up expected calls' responses
-	inPlugin := &csi.GetPluginInfoRequest{}
-	outPlugin := &csi.GetPluginInfoResponse{
-		Name: "foo/bar",
-	}
-	idServer.EXPECT().GetPluginInfo(gomock.Any(), inPlugin).Return(outPlugin, injectedErr).Times(1)
-
 	inProbe := &csi.ProbeRequest{}
 	outProbe := &csi.ProbeResponse{}
 	idServer.EXPECT().Probe(gomock.Any(), inProbe).Return(outProbe, injectedErr).Times(1)
-	// Calling Probing function
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	if err := runProbe(ctx, csiConn); err != nil {
-		t.Fatalf("failed to run probe with error: %+v", err)
+
+	hp := &healthProbe{
+		conn:       csiConn,
+		driverName: driverName,
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if req.URL.String() == "/healthz" {
+			hp.checkProbe(rw, req)
+		}
+	}))
+	defer server.Close()
+
+	httpreq, err := http.NewRequest("GET", server.URL+"/healthz", nil)
+	if err != nil {
+		t.Fatalf("failed to build test request for health check: %v", err)
+	}
+	_, err = http.DefaultClient.Do(httpreq)
+	if err != nil {
+		t.Errorf("failed to check probe: %v", err)
 	}
 }

cmd/main.go

@@ -25,62 +25,47 @@ import (
 
 	"k8s.io/klog"
 
-	"github.com/kubernetes-csi/livenessprobe/pkg/connection"
-)
-
-const (
-	// Default timeout of short CSI calls like GetPluginInfo
-	csiTimeout = time.Second
+	connlib "github.com/kubernetes-csi/csi-lib-utils/connection"
+	"google.golang.org/grpc"
 )
 
 // 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", 0, "The --connection-timeout flag is deprecated")
+	probeTimeout      = flag.Duration("probe-timeout", time.Second, "Probe timeout in seconds")
 	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")
 )
 
-func runProbe(ctx context.Context, csiConn connection.CSIConnection) error {
-	// Get CSI driver name.
-	klog.Infof("Calling CSI driver to discover driver name.")
-	csiDriverName, err := csiConn.GetDriverName(ctx)
-	if err != nil {
-		return err
-	}
-	klog.Infof("CSI driver name: %q", csiDriverName)
-	// Sending Probe request
-	klog.Infof("Sending probe request to CSI driver.")
-	err = csiConn.LivenessProbe(ctx)
-	return err
+type healthProbe struct {
+	conn       *grpc.ClientConn
+	driverName string
 }
 
-func getCSIConnection() (connection.CSIConnection, error) {
-	klog.Infof("Attempting to open a gRPC connection with: %s", *csiAddress)
-	return connection.NewConnection(*csiAddress)
-}
+func (h *healthProbe) checkProbe(w http.ResponseWriter, req *http.Request) {
+	ctx, cancel := context.WithTimeout(req.Context(), *probeTimeout)
+	defer cancel()
 
-func checkHealth(w http.ResponseWriter, req *http.Request) {
-	klog.Infof("Request: %s from: %s\n", req.URL.Path, req.RemoteAddr)
-	csiConn, err := getCSIConnection()
+	klog.Infof("Sending probe request to CSI driver %q", h.driverName)
+	ready, err := connlib.Probe(ctx, h.conn)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		w.Write([]byte(err.Error()))
-		klog.Errorf("Failed to get connection to CSI  with error: %v.", err)
+		klog.Errorf("health check failed: %v", err)
 		return
 	}
-	defer csiConn.Close()
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	if err := runProbe(ctx, csiConn); err != nil {
+
+	if !ready {
 		w.WriteHeader(http.StatusInternalServerError)
 		w.Write([]byte(err.Error()))
-		klog.Errorf("Health check failed with: %v.", err)
-	} else {
-		w.WriteHeader(http.StatusOK)
-		w.Write([]byte(`ok`))
-		klog.Infof("Health check succeeded.")
+		klog.Error("driver responded but is not ready")
+		return
 	}
+
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte(`ok`))
+	klog.Infof("Health check succeeded")
 }
 
 func main() {
@@ -92,10 +77,29 @@ func main() {
 		klog.Warning("--connection-timeout is deprecated and will have no effect")
 	}
 
+	csiConn, err := connlib.Connect(*csiAddress)
+	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
+		klog.Fatalf("failed to establish connection to CSI driver: %v", err)
+	}
+
+	klog.Infof("calling CSI driver to discover driver name")
+	csiDriverName, err := connlib.GetDriverName(context.Background(), csiConn)
+	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,
+	}
+
 	addr := net.JoinHostPort("0.0.0.0", *healthzPort)
-	http.HandleFunc("/healthz", checkHealth)
+	http.HandleFunc("/healthz", hp.checkProbe)
 	klog.Infof("Serving requests to /healthz on: %s", addr)
-	err := http.ListenAndServe(addr, nil)
+	err = http.ListenAndServe(addr, nil)
 	if err != nil {
 		klog.Fatalf("failed to start http server with error: %v", err)
 	}