From bdc484e415793958f10aed93ca418222924deb37 Mon Sep 17 00:00:00 2001
From: Shubheksha Jalan <jshubheksha@gmail.com>
Date: Fri, 25 May 2018 21:08:46 +0530
Subject: [PATCH] Add a timeout to attach and detach calls

---
 cmd/csi-attacher/main.go           |  3 ++-
 pkg/controller/csi_handler.go      | 13 ++++++++++---
 pkg/controller/csi_handler_test.go |  7 ++++++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/cmd/csi-attacher/main.go b/cmd/csi-attacher/main.go
index 011c22485..2a3edb68b 100644
--- a/cmd/csi-attacher/main.go
+++ b/cmd/csi-attacher/main.go
@@ -53,6 +53,7 @@ var (
 	csiAddress        = flag.String("csi-address", "/run/csi/socket", "Address of the CSI driver socket.")
 	dummy             = flag.Bool("dummy", false, "Run in dummy mode, i.e. not connecting to CSI driver and marking everything as attached. Expected CSI driver name is \"csi/dummy\".")
 	showVersion       = flag.Bool("version", false, "Show version.")
+	timeout           = flag.Duration("timeout", 15*time.Second, "Timeout for waiting for attaching or detaching the volume.")
 
 	enableLeaderElection    = flag.Bool("leader-election", false, "Enable leader election.")
 	leaderElectionNamespace = flag.String("leader-election-namespace", "", "Namespace where this attacher runs.")
@@ -138,7 +139,7 @@ func main() {
 				pvLister := factory.Core().V1().PersistentVolumes().Lister()
 				nodeLister := factory.Core().V1().Nodes().Lister()
 				vaLister := factory.Storage().V1beta1().VolumeAttachments().Lister()
-				handler = controller.NewCSIHandler(clientset, attacher, csiConn, pvLister, nodeLister, vaLister)
+				handler = controller.NewCSIHandler(clientset, attacher, csiConn, pvLister, nodeLister, vaLister, timeout)
 				glog.V(2).Infof("CSI driver supports ControllerPublishUnpublish, using real CSI handler")
 			} else {
 				handler = controller.NewTrivialHandler(clientset)
diff --git a/pkg/controller/csi_handler.go b/pkg/controller/csi_handler.go
index 351b7da7e..bb8c9d7c4 100644
--- a/pkg/controller/csi_handler.go
+++ b/pkg/controller/csi_handler.go
@@ -19,6 +19,7 @@ package controller
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/golang/glog"
 
@@ -45,17 +46,20 @@ type csiHandler struct {
 	nodeLister       corelisters.NodeLister
 	vaLister         storagelisters.VolumeAttachmentLister
 	vaQueue, pvQueue workqueue.RateLimitingInterface
+	timeout          time.Duration
 }
 
 var _ Handler = &csiHandler{}
 
+// NewCSIHandler creates a new CSIHandler.
 func NewCSIHandler(
 	client kubernetes.Interface,
 	attacherName string,
 	csiConnection connection.CSIConnection,
 	pvLister corelisters.PersistentVolumeLister,
 	nodeLister corelisters.NodeLister,
-	vaLister storagelisters.VolumeAttachmentLister) Handler {
+	vaLister storagelisters.VolumeAttachmentLister,
+	timeout *time.Duration) Handler {
 
 	return &csiHandler{
 		client:        client,
@@ -64,6 +68,7 @@ func NewCSIHandler(
 		pvLister:      pvLister,
 		nodeLister:    nodeLister,
 		vaLister:      vaLister,
+		timeout:       *timeout,
 	}
 }
 
@@ -258,7 +263,8 @@ func (h *csiHandler) csiAttach(va *storage.VolumeAttachment) (*storage.VolumeAtt
 		return va, nil, fmt.Errorf("could not add VolumeAttachment finalizer: %s", err)
 	}
 
-	ctx := context.TODO()
+	ctx, cancel := context.WithTimeout(context.Background(), h.timeout)
+	defer cancel()
 	// We're not interested in `detached` return value, the controller will
 	// issue Detach to be sure the volume is really detached.
 	publishInfo, _, err := h.csiConnection.Attach(ctx, volumeHandle, readOnly, nodeID, volumeCapabilities, attributes, secrets)
@@ -296,7 +302,8 @@ func (h *csiHandler) csiDetach(va *storage.VolumeAttachment) (*storage.VolumeAtt
 		return va, err
 	}
 
-	ctx := context.TODO()
+	ctx, cancel := context.WithTimeout(context.Background(), h.timeout)
+	defer cancel()
 	detached, err := h.csiConnection.Detach(ctx, volumeHandle, nodeID, secrets)
 	if err != nil && !detached {
 		// The volume may not be fully detached. Save the error and try again
diff --git a/pkg/controller/csi_handler_test.go b/pkg/controller/csi_handler_test.go
index 455cf2f00..55d8f7866 100644
--- a/pkg/controller/csi_handler_test.go
+++ b/pkg/controller/csi_handler_test.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/kubernetes-csi/external-attacher/pkg/connection"
 
@@ -39,6 +40,8 @@ const (
 	fin = "external-attacher/csi-test"
 )
 
+var timeout = 1 * time.Minute
+
 func csiHandlerFactory(client kubernetes.Interface, informerFactory informers.SharedInformerFactory, csi connection.CSIConnection) Handler {
 	return NewCSIHandler(
 		client,
@@ -46,7 +49,9 @@ func csiHandlerFactory(client kubernetes.Interface, informerFactory informers.Sh
 		csi,
 		informerFactory.Core().V1().PersistentVolumes().Lister(),
 		informerFactory.Core().V1().Nodes().Lister(),
-		informerFactory.Storage().V1beta1().VolumeAttachments().Lister())
+		informerFactory.Storage().V1beta1().VolumeAttachments().Lister(),
+		&timeout,
+	)
 }
 
 func pv() *v1.PersistentVolume {