interop: Replace context.Background() with passed ctx (#6827)

Author: Aditya-Sood

interop/client/client.go

@@ -169,6 +169,8 @@ func main() {
 		logger.Fatalf("only one of TLS, ALTS, google default creds, or compute engine creds can be used")
 	}
 
+	ctx := context.Background()
+
 	var credsChosen credsMode
 	switch {
 	case *useTLS:
@@ -242,7 +244,7 @@ func main() {
 			}
 			opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds))
 		} else if *testCase == "oauth2_auth_token" {
-			opts = append(opts, grpc.WithPerRPCCredentials(oauth.TokenSource{TokenSource: oauth2.StaticTokenSource(interop.GetToken(*serviceAccountKeyFile, *oauthScope))}))
+			opts = append(opts, grpc.WithPerRPCCredentials(oauth.TokenSource{TokenSource: oauth2.StaticTokenSource(interop.GetToken(ctx, *serviceAccountKeyFile, *oauthScope))}))
 		}
 	}
 	if len(*serviceConfigJSON) > 0 {
@@ -265,105 +267,107 @@ func main() {
 	}
 	defer conn.Close()
 	tc := testgrpc.NewTestServiceClient(conn)
+	ctxWithDeadline, cancel := context.WithTimeout(ctx, time.Duration(*soakOverallTimeoutSeconds)*time.Second)
+	defer cancel()
 	switch *testCase {
 	case "empty_unary":
-		interop.DoEmptyUnaryCall(tc)
+		interop.DoEmptyUnaryCall(ctx, tc)
 		logger.Infoln("EmptyUnaryCall done")
 	case "large_unary":
-		interop.DoLargeUnaryCall(tc)
+		interop.DoLargeUnaryCall(ctx, tc)
 		logger.Infoln("LargeUnaryCall done")
 	case "client_streaming":
-		interop.DoClientStreaming(tc)
+		interop.DoClientStreaming(ctx, tc)
 		logger.Infoln("ClientStreaming done")
 	case "server_streaming":
-		interop.DoServerStreaming(tc)
+		interop.DoServerStreaming(ctx, tc)
 		logger.Infoln("ServerStreaming done")
 	case "ping_pong":
-		interop.DoPingPong(tc)
+		interop.DoPingPong(ctx, tc)
 		logger.Infoln("Pingpong done")
 	case "empty_stream":
-		interop.DoEmptyStream(tc)
+		interop.DoEmptyStream(ctx, tc)
 		logger.Infoln("Emptystream done")
 	case "timeout_on_sleeping_server":
-		interop.DoTimeoutOnSleepingServer(tc)
+		interop.DoTimeoutOnSleepingServer(ctx, tc)
 		logger.Infoln("TimeoutOnSleepingServer done")
 	case "compute_engine_creds":
 		if credsChosen != credsTLS {
 			logger.Fatalf("TLS credentials need to be set for compute_engine_creds test case.")
 		}
-		interop.DoComputeEngineCreds(tc, *defaultServiceAccount, *oauthScope)
+		interop.DoComputeEngineCreds(ctx, tc, *defaultServiceAccount, *oauthScope)
 		logger.Infoln("ComputeEngineCreds done")
 	case "service_account_creds":
 		if credsChosen != credsTLS {
 			logger.Fatalf("TLS credentials need to be set for service_account_creds test case.")
 		}
-		interop.DoServiceAccountCreds(tc, *serviceAccountKeyFile, *oauthScope)
+		interop.DoServiceAccountCreds(ctx, tc, *serviceAccountKeyFile, *oauthScope)
 		logger.Infoln("ServiceAccountCreds done")
 	case "jwt_token_creds":
 		if credsChosen != credsTLS {
 			logger.Fatalf("TLS credentials need to be set for jwt_token_creds test case.")
 		}
-		interop.DoJWTTokenCreds(tc, *serviceAccountKeyFile)
+		interop.DoJWTTokenCreds(ctx, tc, *serviceAccountKeyFile)
 		logger.Infoln("JWTtokenCreds done")
 	case "per_rpc_creds":
 		if credsChosen != credsTLS {
 			logger.Fatalf("TLS credentials need to be set for per_rpc_creds test case.")
 		}
-		interop.DoPerRPCCreds(tc, *serviceAccountKeyFile, *oauthScope)
+		interop.DoPerRPCCreds(ctx, tc, *serviceAccountKeyFile, *oauthScope)
 		logger.Infoln("PerRPCCreds done")
 	case "oauth2_auth_token":
 		if credsChosen != credsTLS {
 			logger.Fatalf("TLS credentials need to be set for oauth2_auth_token test case.")
 		}
-		interop.DoOauth2TokenCreds(tc, *serviceAccountKeyFile, *oauthScope)
+		interop.DoOauth2TokenCreds(ctx, tc, *serviceAccountKeyFile, *oauthScope)
 		logger.Infoln("Oauth2TokenCreds done")
 	case "google_default_credentials":
 		if credsChosen != credsGoogleDefaultCreds {
 			logger.Fatalf("GoogleDefaultCredentials need to be set for google_default_credentials test case.")
 		}
-		interop.DoGoogleDefaultCredentials(tc, *defaultServiceAccount)
+		interop.DoGoogleDefaultCredentials(ctx, tc, *defaultServiceAccount)
 		logger.Infoln("GoogleDefaultCredentials done")
 	case "compute_engine_channel_credentials":
 		if credsChosen != credsComputeEngineCreds {
 			logger.Fatalf("ComputeEngineCreds need to be set for compute_engine_channel_credentials test case.")
 		}
-		interop.DoComputeEngineChannelCredentials(tc, *defaultServiceAccount)
+		interop.DoComputeEngineChannelCredentials(ctx, tc, *defaultServiceAccount)
 		logger.Infoln("ComputeEngineChannelCredentials done")
 	case "cancel_after_begin":
-		interop.DoCancelAfterBegin(tc)
+		interop.DoCancelAfterBegin(ctx, tc)
 		logger.Infoln("CancelAfterBegin done")
 	case "cancel_after_first_response":
-		interop.DoCancelAfterFirstResponse(tc)
+		interop.DoCancelAfterFirstResponse(ctx, tc)
 		logger.Infoln("CancelAfterFirstResponse done")
 	case "status_code_and_message":
-		interop.DoStatusCodeAndMessage(tc)
+		interop.DoStatusCodeAndMessage(ctx, tc)
 		logger.Infoln("StatusCodeAndMessage done")
 	case "special_status_message":
-		interop.DoSpecialStatusMessage(tc)
+		interop.DoSpecialStatusMessage(ctx, tc)
 		logger.Infoln("SpecialStatusMessage done")
 	case "custom_metadata":
-		interop.DoCustomMetadata(tc)
+		interop.DoCustomMetadata(ctx, tc)
 		logger.Infoln("CustomMetadata done")
 	case "unimplemented_method":
-		interop.DoUnimplementedMethod(conn)
+		interop.DoUnimplementedMethod(conn, ctx)
 		logger.Infoln("UnimplementedMethod done")
 	case "unimplemented_service":
-		interop.DoUnimplementedService(testgrpc.NewUnimplementedServiceClient(conn))
+		interop.DoUnimplementedService(testgrpc.NewUnimplementedServiceClient(conn), ctx)
 		logger.Infoln("UnimplementedService done")
 	case "pick_first_unary":
-		interop.DoPickFirstUnary(tc)
+		interop.DoPickFirstUnary(ctx, tc)
 		logger.Infoln("PickFirstUnary done")
 	case "rpc_soak":
-		interop.DoSoakTest(tc, serverAddr, opts, false /* resetChannel */, *soakIterations, *soakMaxFailures, *soakRequestSize, *soakResponseSize, time.Duration(*soakPerIterationMaxAcceptableLatencyMs)*time.Millisecond, time.Duration(*soakMinTimeMsBetweenRPCs)*time.Millisecond, time.Now().Add(time.Duration(*soakOverallTimeoutSeconds)*time.Second))
+		interop.DoSoakTest(ctxWithDeadline, tc, serverAddr, opts, false /* resetChannel */, *soakIterations, *soakMaxFailures, *soakRequestSize, *soakResponseSize, time.Duration(*soakPerIterationMaxAcceptableLatencyMs)*time.Millisecond, time.Duration(*soakMinTimeMsBetweenRPCs)*time.Millisecond)
 		logger.Infoln("RpcSoak done")
 	case "channel_soak":
-		interop.DoSoakTest(tc, serverAddr, opts, true /* resetChannel */, *soakIterations, *soakMaxFailures, *soakRequestSize, *soakResponseSize, time.Duration(*soakPerIterationMaxAcceptableLatencyMs)*time.Millisecond, time.Duration(*soakMinTimeMsBetweenRPCs)*time.Millisecond, time.Now().Add(time.Duration(*soakOverallTimeoutSeconds)*time.Second))
+		interop.DoSoakTest(ctxWithDeadline, tc, serverAddr, opts, true /* resetChannel */, *soakIterations, *soakMaxFailures, *soakRequestSize, *soakResponseSize, time.Duration(*soakPerIterationMaxAcceptableLatencyMs)*time.Millisecond, time.Duration(*soakMinTimeMsBetweenRPCs)*time.Millisecond)
 		logger.Infoln("ChannelSoak done")
 	case "orca_per_rpc":
-		interop.DoORCAPerRPCTest(tc)
+		interop.DoORCAPerRPCTest(ctx, tc)
 		logger.Infoln("ORCAPerRPC done")
 	case "orca_oob":
-		interop.DoORCAOOBTest(tc)
+		interop.DoORCAOOBTest(ctx, tc)
 		logger.Infoln("ORCAOOB done")
 	default:
 		logger.Fatal("Unsupported test case: ", *testCase)

interop/http2/negative_http2_client.go

@@ -69,12 +69,12 @@ func largeSimpleRequest() *testpb.SimpleRequest {
 }
 
 // sends two unary calls. The server asserts that the calls use different connections.
-func goaway(tc testgrpc.TestServiceClient) {
-	interop.DoLargeUnaryCall(tc)
+func goaway(ctx context.Context, tc testgrpc.TestServiceClient) {
+	interop.DoLargeUnaryCall(ctx, tc)
 	// sleep to ensure that the client has time to recv the GOAWAY.
 	// TODO(ncteisen): make this less hacky.
 	time.Sleep(1 * time.Second)
-	interop.DoLargeUnaryCall(tc)
+	interop.DoLargeUnaryCall(ctx, tc)
 }
 
 func rstAfterHeader(tc testgrpc.TestServiceClient) {
@@ -110,19 +110,19 @@ func rstAfterData(tc testgrpc.TestServiceClient) {
 	}
 }
 
-func ping(tc testgrpc.TestServiceClient) {
+func ping(ctx context.Context, tc testgrpc.TestServiceClient) {
 	// The server will assert that every ping it sends was ACK-ed by the client.
-	interop.DoLargeUnaryCall(tc)
+	interop.DoLargeUnaryCall(ctx, tc)
 }
 
-func maxStreams(tc testgrpc.TestServiceClient) {
-	interop.DoLargeUnaryCall(tc)
+func maxStreams(ctx context.Context, tc testgrpc.TestServiceClient) {
+	interop.DoLargeUnaryCall(ctx, tc)
 	var wg sync.WaitGroup
 	for i := 0; i < 15; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			interop.DoLargeUnaryCall(tc)
+			interop.DoLargeUnaryCall(ctx, tc)
 		}()
 	}
 	wg.Wait()
@@ -139,9 +139,10 @@ func main() {
 	}
 	defer conn.Close()
 	tc := testgrpc.NewTestServiceClient(conn)
+	ctx := context.Background()
 	switch *testCase {
 	case "goaway":
-		goaway(tc)
+		goaway(ctx, tc)
 		logger.Infoln("goaway done")
 	case "rst_after_header":
 		rstAfterHeader(tc)
@@ -153,10 +154,10 @@ func main() {
 		rstAfterData(tc)
 		logger.Infoln("rst_after_data done")
 	case "ping":
-		ping(tc)
+		ping(ctx, tc)
 		logger.Infoln("ping done")
 	case "max_streams":
-		maxStreams(tc)
+		maxStreams(ctx, tc)
 		logger.Infoln("max_streams done")
 	default:
 		logger.Fatal("Unsupported test case: ", *testCase)

interop/observability/client/client.go

@@ -58,13 +58,14 @@ func main() {
 	}
 	defer conn.Close()
 	tc := testgrpc.NewTestServiceClient(conn)
+	ctx := context.Background()
 	for i := 0; i < *numTimes; i++ {
 		if *testCase == "ping_pong" {
-			interop.DoPingPong(tc)
+			interop.DoPingPong(ctx, tc)
 		} else if *testCase == "large_unary" {
-			interop.DoLargeUnaryCall(tc)
+			interop.DoLargeUnaryCall(ctx, tc)
 		} else if *testCase == "custom_metadata" {
-			interop.DoCustomMetadata(tc)
+			interop.DoCustomMetadata(ctx, tc)
 		} else {
 			log.Fatalf("Invalid test case: %s", *testCase)
 		}

interop/stress/client/main.go

@@ -226,32 +226,33 @@ func startServer(server *server, port int) {
 func performRPCs(gauge *gauge, conn *grpc.ClientConn, selector *weightedRandomTestSelector, stop <-chan bool) {
 	client := testgrpc.NewTestServiceClient(conn)
 	var numCalls int64
+	ctx := context.Background()
 	startTime := time.Now()
 	for {
 		test := selector.getNextTest()
 		switch test {
 		case "empty_unary":
-			interop.DoEmptyUnaryCall(client)
+			interop.DoEmptyUnaryCall(ctx, client)
 		case "large_unary":
-			interop.DoLargeUnaryCall(client)
+			interop.DoLargeUnaryCall(ctx, client)
 		case "client_streaming":
-			interop.DoClientStreaming(client)
+			interop.DoClientStreaming(ctx, client)
 		case "server_streaming":
-			interop.DoServerStreaming(client)
+			interop.DoServerStreaming(ctx, client)
 		case "ping_pong":
-			interop.DoPingPong(client)
+			interop.DoPingPong(ctx, client)
 		case "empty_stream":
-			interop.DoEmptyStream(client)
+			interop.DoEmptyStream(ctx, client)
 		case "timeout_on_sleeping_server":
-			interop.DoTimeoutOnSleepingServer(client)
+			interop.DoTimeoutOnSleepingServer(ctx, client)
 		case "cancel_after_begin":
-			interop.DoCancelAfterBegin(client)
+			interop.DoCancelAfterBegin(ctx, client)
 		case "cancel_after_first_response":
-			interop.DoCancelAfterFirstResponse(client)
+			interop.DoCancelAfterFirstResponse(ctx, client)
 		case "status_code_and_message":
-			interop.DoStatusCodeAndMessage(client)
+			interop.DoStatusCodeAndMessage(ctx, client)
 		case "custom_metadata":
-			interop.DoCustomMetadata(client)
+			interop.DoCustomMetadata(ctx, client)
 		}
 		numCalls++
 		defer func() { atomic.AddInt64(&totalNumCalls, numCalls) }()