Skip to content

Commit

Permalink
interop: Replace context.Background() with passed ctx (#6827)
Browse files Browse the repository at this point in the history
  • Loading branch information
Aditya-Sood authored Jan 31, 2024
1 parent 3aafa84 commit a3f5ed6
Show file tree
Hide file tree
Showing 6 changed files with 121 additions and 110 deletions.
58 changes: 31 additions & 27 deletions interop/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
Expand Down
23 changes: 12 additions & 11 deletions interop/http2/negative_http2_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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()
Expand All @@ -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)
Expand All @@ -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)
Expand Down
7 changes: 4 additions & 3 deletions interop/observability/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
23 changes: 12 additions & 11 deletions interop/stress/client/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) }()
Expand Down
Loading

0 comments on commit a3f5ed6

Please sign in to comment.