Merge pull request #16 from jsafrane/probe

Add ControllerProbe

cmd/csi-attacher/main.go

@@ -19,6 +19,7 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"os"
 	"os/signal"
 	"time"
@@ -96,6 +97,12 @@ func main() {
 		}
 		glog.V(2).Infof("CSI driver name: %q", attacher)
 
+		// Check it's ready
+		if err = waitForDriverReady(csiConn, *connectionTimeout); err != nil {
+			glog.Error(err.Error())
+			os.Exit(1)
+		}
+
 		// Find out if the driver supports attach/detach.
 		supportsAttach, err := csiConn.SupportsControllerPublish(ctx)
 		if err != nil {
@@ -140,3 +147,26 @@ func buildConfig(kubeconfig string) (*rest.Config, error) {
 	}
 	return rest.InClusterConfig()
 }
+
+func waitForDriverReady(csiConn connection.CSIConnection, timeout time.Duration) error {
+	now := time.Now()
+	finish := now.Add(timeout)
+	var err error
+
+	for {
+		ctx, cancel := context.WithTimeout(context.Background(), csiTimeout)
+		defer cancel()
+		err = csiConn.ControllerProbe(ctx)
+		if err == nil {
+			glog.V(2).Infof("ControllerProbe succeeded")
+			return nil
+		}
+		glog.V(2).Infof("ControllerProbe failed with %s", err)
+
+		now := time.Now()
+		if now.After(finish) {
+			return fmt.Errorf("Failed to probe the controller: %s", err)
+		}
+		time.Sleep(time.Second)
+	}
+}

pkg/connection/connection.go

@@ -55,6 +55,9 @@ type CSIConnection interface {
 	// should retry.
 	Detach(ctx context.Context, volumeID string, nodeID string) (detached bool, err error)
 
+	// Probe checks that the CSI driver is ready to process requests
+	ControllerProbe(ctx context.Context) error
+
 	// Close the connection
 	Close() error
 }
@@ -197,6 +200,17 @@ func (c *csiConnection) Detach(ctx context.Context, volumeID string, nodeID stri
 	return true, nil
 }
 
+func (c *csiConnection) ControllerProbe(ctx context.Context) error {
+	client := csi.NewControllerClient(c.conn)
+
+	req := csi.ControllerProbeRequest{
+		Version: &csiVersion,
+	}
+
+	_, err := client.ControllerProbe(ctx, &req)
+	return err
+}
+
 func (c *csiConnection) Close() error {
 	return c.conn.Close()
 }

pkg/connection/connection_test.go

@@ -485,3 +485,55 @@ func TestDetachAttach(t *testing.T) {
 		}
 	}
 }
+
+func TestControllerProbe(t *testing.T) {
+	tests := []struct {
+		name        string
+		injectError bool
+		expectError bool
+	}{
+		{
+			name:        "success",
+			expectError: false,
+		},
+		{
+			name:        "gRPC error",
+			injectError: true,
+			expectError: true,
+		},
+	}
+
+	mockController, driver, _, controllerServer, csiConn, err := createMockServer(t)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer mockController.Finish()
+	defer driver.Stop()
+	defer csiConn.Close()
+
+	for _, test := range tests {
+		in := &csi.ControllerProbeRequest{
+			Version: &csi.Version{
+				Major: 0,
+				Minor: 1,
+				Patch: 0,
+			},
+		}
+		out := &csi.ControllerProbeResponse{}
+		var injectedErr error = nil
+		if test.injectError {
+			injectedErr = fmt.Errorf("mock error")
+		}
+
+		// Setup expectation
+		controllerServer.EXPECT().ControllerProbe(gomock.Any(), in).Return(out, injectedErr).Times(1)
+
+		err := csiConn.ControllerProbe(context.Background())
+		if test.expectError && err == nil {
+			t.Errorf("test %q: Expected error, got none", test.name)
+		}
+		if !test.expectError && err != nil {
+			t.Errorf("test %q: got error: %v", test.name, err)
+		}
+	}
+}

pkg/controller/framework_test.go

@@ -394,3 +394,7 @@ func (f *fakeCSIConnection) Detach(ctx context.Context, volumeID string, nodeID
 func (f *fakeCSIConnection) Close() error {
 	return fmt.Errorf("Not implemented")
 }
+
+func (f *fakeCSIConnection) ControllerProbe(ctx context.Context) error {
+	return fmt.Errorf("Not implemented")
+}