Thread contexts into all appropriate functions in non-test code

Author: davidz627

cmd/main.go

@@ -15,6 +15,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"math/rand"
 	"os"
@@ -63,8 +64,11 @@ func handle() {
 
 	gceDriver := driver.GetGCEDriver()
 
-	//Initialize GCE Driver (Move setup to main?)
-	cloudProvider, err := gce.CreateCloudProvider(vendorVersion, *gceConfigFilePath)
+	//Initialize GCE Driver
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cloudProvider, err := gce.CreateCloudProvider(ctx, vendorVersion, *gceConfigFilePath)
 	if err != nil {
 		klog.Fatalf("Failed to get cloud provider: %v", err)
 	}

pkg/gce-cloud-provider/compute/gce.go

@@ -68,7 +68,7 @@ type ConfigGlobal struct {
 	ProjectId string `gcfg:"project-id"`
 }
 
-func CreateCloudProvider(vendorVersion string, configPath string) (*CloudProvider, error) {
+func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string) (*CloudProvider, error) {
 	configFile, err := readConfig(configPath)
 	if err != nil {
 		return nil, err
@@ -78,12 +78,12 @@ func CreateCloudProvider(vendorVersion string, configPath string) (*CloudProvide
 
 	klog.V(2).Infof("Using GCE provider config %+v", configFile)
 
-	tokenSource, err := generateTokenSource(configFile)
+	tokenSource, err := generateTokenSource(ctx, configFile)
 	if err != nil {
 		return nil, err
 	}
 
-	svc, err := createCloudService(vendorVersion, tokenSource)
+	svc, err := createCloudService(ctx, vendorVersion, tokenSource)
 	if err != nil {
 		return nil, err
 	}
@@ -102,7 +102,7 @@ func CreateCloudProvider(vendorVersion string, configPath string) (*CloudProvide
 
 }
 
-func generateTokenSource(configFile *ConfigFile) (oauth2.TokenSource, error) {
+func generateTokenSource(ctx context.Context, configFile *ConfigFile) (oauth2.TokenSource, error) {
 
 	if configFile != nil && configFile.Global.TokenURL != "" && configFile.Global.TokenURL != "nil" {
 		// configFile.Global.TokenURL is defined
@@ -116,7 +116,7 @@ func generateTokenSource(configFile *ConfigFile) (oauth2.TokenSource, error) {
 	// Use DefaultTokenSource
 
 	tokenSource, err := google.DefaultTokenSource(
-		context.Background(),
+		ctx,
 		compute.CloudPlatformScope,
 		compute.ComputeScope)
 
@@ -149,13 +149,13 @@ func readConfig(configPath string) (*ConfigFile, error) {
 	return cfg, nil
 }
 
-func createCloudService(vendorVersion string, tokenSource oauth2.TokenSource) (*compute.Service, error) {
-	svc, err := createCloudServiceWithDefaultServiceAccount(vendorVersion, tokenSource)
+func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource) (*compute.Service, error) {
+	svc, err := createCloudServiceWithDefaultServiceAccount(ctx, vendorVersion, tokenSource)
 	return svc, err
 }
 
-func createCloudServiceWithDefaultServiceAccount(vendorVersion string, tokenSource oauth2.TokenSource) (*compute.Service, error) {
-	client, err := newOauthClient(tokenSource)
+func createCloudServiceWithDefaultServiceAccount(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource) (*compute.Service, error) {
+	client, err := newOauthClient(ctx, tokenSource)
 	if err != nil {
 		return nil, err
 	}
@@ -167,7 +167,7 @@ func createCloudServiceWithDefaultServiceAccount(vendorVersion string, tokenSour
 	return service, nil
 }
 
-func newOauthClient(tokenSource oauth2.TokenSource) (*http.Client, error) {
+func newOauthClient(ctx context.Context, tokenSource oauth2.TokenSource) (*http.Client, error) {
 	if err := wait.PollImmediate(5*time.Second, 30*time.Second, func() (bool, error) {
 		if _, err := tokenSource.Token(); err != nil {
 			klog.Errorf("error fetching initial token: %v", err)
@@ -178,7 +178,7 @@ func newOauthClient(tokenSource oauth2.TokenSource) (*http.Client, error) {
 		return nil, err
 	}
 
-	return oauth2.NewClient(context.Background(), tokenSource), nil
+	return oauth2.NewClient(ctx, tokenSource), nil
 }
 
 func getProjectAndZone(config *ConfigFile) (string, string, error) {

pkg/gce-pd-csi-driver/controller.go

@@ -118,7 +118,7 @@ func (gceCS *GCEControllerServer) CreateVolume(ctx context.Context, req *csi.Cre
 	var volKey *meta.Key
 	switch replicationType {
 	case replicationTypeNone:
-		zones, err = pickZones(gceCS, req.GetAccessibilityRequirements(), 1)
+		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 1)
 		if err != nil {
 			return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("CreateVolume failed to pick zones for disk: %v", err))
 		}
@@ -128,7 +128,7 @@ func (gceCS *GCEControllerServer) CreateVolume(ctx context.Context, req *csi.Cre
 		volKey = meta.ZonalKey(name, zones[0])
 
 	case replicationTypeRegionalPD:
-		zones, err = pickZones(gceCS, req.GetAccessibilityRequirements(), 2)
+		zones, err = pickZones(ctx, gceCS, req.GetAccessibilityRequirements(), 2)
 		if err != nil {
 			return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("CreateVolume failed to pick zones for disk: %v", err))
 		}
@@ -941,7 +941,7 @@ func getZoneFromSegment(seg map[string]string) (string, error) {
 	return zone, nil
 }
 
-func pickZones(gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZones int) ([]string, error) {
+func pickZones(ctx context.Context, gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZones int) ([]string, error) {
 	var zones []string
 	var err error
 	if top != nil {
@@ -950,7 +950,7 @@ func pickZones(gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZone
 			return nil, fmt.Errorf("failed to pick zones from topology: %v", err)
 		}
 	} else {
-		zones, err = getDefaultZonesInRegion(gceCS, []string{gceCS.MetadataService.GetZone()}, numZones)
+		zones, err = getDefaultZonesInRegion(ctx, gceCS, []string{gceCS.MetadataService.GetZone()}, numZones)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get default %v zones in region: %v", numZones, err)
 		}
@@ -960,13 +960,13 @@ func pickZones(gceCS *GCEControllerServer, top *csi.TopologyRequirement, numZone
 	return zones, nil
 }
 
-func getDefaultZonesInRegion(gceCS *GCEControllerServer, existingZones []string, numZones int) ([]string, error) {
+func getDefaultZonesInRegion(ctx context.Context, gceCS *GCEControllerServer, existingZones []string, numZones int) ([]string, error) {
 	region, err := common.GetRegionFromZones(existingZones)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get region from zones: %v", err)
 	}
 	needToGet := numZones - len(existingZones)
-	totZones, err := gceCS.CloudProvider.ListZones(context.Background(), region)
+	totZones, err := gceCS.CloudProvider.ListZones(ctx, region)
 	if err != nil {
 		return nil, fmt.Errorf("failed to list zones from cloud provider: %v", err)
 	}