Add option to exit on reconnect

connection/connection.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"strings"
 	"time"
@@ -33,6 +34,9 @@ import (
 const (
 	// Interval of logging connection errors
 	connectionLoggingInterval = 10 * time.Second
+
+	// Default path to termination log
+	TerminationLogPath = "/dev/termination-log"
 )
 
 // Connect opens insecure gRPC connection to a CSI driver. Address must be either absolute path to UNIX domain socket
@@ -65,12 +69,32 @@ type Option func(o *options)
 // connection got lost. If that callback returns true, the connection
 // is restablished. Otherwise the connection is left as it is and
 // all future gRPC calls using it will fail with status.Unavailable.
+// Only one of ExitOnConnectionLoss and OnConnectionLoss can be used.
 func OnConnectionLoss(reconnect func() bool) Option {
 	return func(o *options) {
 		o.reconnect = reconnect
 	}
 }
 
+// ExitOnConnectionLoss exits when connection gets lost. Optionally it
+// writes error to given file. Use with /dev/termination-log to get a nice
+// error in Kubernetes.
+// Only one of ExitOnConnectionLoss and OnConnectionLoss can be used.
+func ExitOnConnectionLoss(terminationLogPath string) Option {
+	return func(o *options) {
+		o.reconnect = func() bool {
+			terminationMsg := "Lost connection to CSI driver, exiting"
+			if terminationLogPath != "" {
+				if err := ioutil.WriteFile(terminationLogPath, []byte(terminationMsg), 0644); err != nil {
+					klog.Errorf("%s: %s", terminationLogPath, err)
+				}
+			}
+			klog.Fatalf(terminationMsg)
+			return false
+		}
+	}
+}
+
 type options struct {
 	reconnect func() bool
 }