refactor external snapshotter to use csi-lib-utils/rpc

cmd/csi-snapshotter/main.go

@@ -24,13 +24,17 @@ import (
 	"os/signal"
 	"time"
 
-	"github.com/golang/glog"
+	"google.golang.org/grpc"
+
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 
-	"github.com/kubernetes-csi/external-snapshotter/pkg/connection"
+	"github.com/container-storage-interface/spec/lib/go/csi"
+	"github.com/golang/glog"
+	"github.com/kubernetes-csi/csi-lib-utils/connection"
+	csirpc "github.com/kubernetes-csi/csi-lib-utils/rpc"
 	"github.com/kubernetes-csi/external-snapshotter/pkg/controller"
 
 	clientset "github.com/kubernetes-csi/external-snapshotter/pkg/client/clientset/versioned"
@@ -123,9 +127,9 @@ func main() {
 	snapshotscheme.AddToScheme(scheme.Scheme)
 
 	// Connect to CSI.
-	csiConn, err := connection.New(*csiAddress)
+	csiConn, err := connection.Connect(*csiAddress)
 	if err != nil {
-		glog.Error(err.Error())
+		glog.Errorf("error connecting to CSI driver: %v", err)
 		os.Exit(1)
 	}
 
@@ -134,23 +138,25 @@ func main() {
 	defer cancel()
 
 	// Find driver name
-	*snapshotter, err = csiConn.GetDriverName(ctx)
+	*snapshotter, err = csirpc.GetDriverName(ctx, csiConn)
 	if err != nil {
-		glog.Error(err.Error())
+		glog.Errorf("error getting CSI driver name: %v", err)
 		os.Exit(1)
 	}
+
 	glog.V(2).Infof("CSI driver name: %q", *snapshotter)
 
 	// Check it's ready
-	if err = waitForDriverReady(csiConn, *connectionTimeout); err != nil {
-		glog.Error(err.Error())
+	if err = csirpc.ProbeForever(csiConn, csiTimeout); err != nil {
+		glog.Errorf("error waiting for CSI driver to be ready: %v", err)
 		os.Exit(1)
+
 	}
 
 	// Find out if the driver supports create/delete snapshot.
-	supportsCreateSnapshot, err := csiConn.SupportsControllerCreateSnapshot(ctx)
+	supportsCreateSnapshot, err := supportsControllerCreateSnapshot(ctx, csiConn)
 	if err != nil {
-		glog.Error(err.Error())
+		glog.Errorf("error determining if driver supports create/delete snapshot operations: %v", err)
 		os.Exit(1)
 	}
 	if !supportsCreateSnapshot {
@@ -165,6 +171,7 @@ func main() {
 
 	glog.V(2).Infof("Start NewCSISnapshotController with snapshotter [%s] kubeconfig [%s] connectionTimeout [%+v] csiAddress [%s] createSnapshotContentRetryCount [%d] createSnapshotContentInterval [%+v] resyncPeriod [%+v] snapshotNamePrefix [%s] snapshotNameUUIDLength [%d]", *snapshotter, *kubeconfig, *connectionTimeout, *csiAddress, createSnapshotContentRetryCount, *createSnapshotContentInterval, *resyncPeriod, *snapshotNamePrefix, snapshotNameUUIDLength)
 
+	snapShotter := controller.NewSnapshotter(csiConn)
 	ctrl := controller.NewCSISnapshotController(
 		snapClient,
 		kubeClient,
@@ -175,7 +182,7 @@ func main() {
 		coreFactory.Core().V1().PersistentVolumeClaims(),
 		*createSnapshotContentRetryCount,
 		*createSnapshotContentInterval,
-		csiConn,
+		snapShotter,
 		*connectionTimeout,
 		*resyncPeriod,
 		*snapshotNamePrefix,
@@ -202,24 +209,11 @@ 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.Probe(ctx)
-		if err == nil {
-			glog.V(2).Infof("Probe succeeded")
-			return nil
-		}
-		glog.V(2).Infof("Probe 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)
+func supportsControllerCreateSnapshot(ctx context.Context, conn *grpc.ClientConn) (bool, error) {
+	capabilities, err := csirpc.GetControllerCapabilities(ctx, conn)
+	if err != nil {
+		return false, err
 	}
+
+	return capabilities[csi.ControllerServiceCapability_RPC_CREATE_DELETE_SNAPSHOT], nil
 }