Use grpc context for GetSystemInfo. (temporalio#1528)

Quinn-With-Two-Ns · Jun 24, 2024 · 5c6a928 · 5c6a928
1 parent 652b15d
commit 5c6a928
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/internal/grpc_dialer_test.go b/internal/grpc_dialer_test.go
@@ -245,11 +245,33 @@ func TestEagerAndLazyClient(t *testing.T) {
 	srv.getSystemInfoResponseError = nil
 	err = c.SignalWorkflow(context.Background(), "workflow1", "", "my-signal", nil)
 	require.NoError(t, err)
+	// Verify version headers are set
+	require.Equal(
+		t,
+		[]string{SDKVersion},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientVersionHeaderName),
+	)
+	require.Equal(
+		t,
+		[]string{clientNameHeaderValue},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientNameHeaderName),
+	)
 
 	// Now that there's no sys info response error, eager should succeed
 	c, err = DialClient(context.Background(), ClientOptions{HostPort: srv.addr})
 	require.NoError(t, err)
 	defer c.Close()
+	// Verify version headers are set
+	require.Equal(
+		t,
+		[]string{SDKVersion},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientVersionHeaderName),
+	)
+	require.Equal(
+		t,
+		[]string{clientNameHeaderValue},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientNameHeaderName),
+	)
 
 	// And even if it starts erroring, the success was memoized so calls succeed
 	srv.getSystemInfoResponseError = fmt.Errorf("some server failure")
@@ -456,6 +478,17 @@ func TestCredentialsAPIKey(t *testing.T) {
 		[]string{"Bearer my-api-key"},
 		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, "Authorization"),
 	)
+	// Verify version headers are set
+	require.Equal(
+		t,
+		[]string{SDKVersion},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientVersionHeaderName),
+	)
+	require.Equal(
+		t,
+		[]string{clientNameHeaderValue},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientNameHeaderName),
+	)
 
 	// Overwrite via context
 	_, err = client.WorkflowService().GetSystemInfo(
@@ -483,6 +516,18 @@ func TestCredentialsAPIKey(t *testing.T) {
 		[]string{"Bearer my-callback-api-key"},
 		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, "Authorization"),
 	)
+
+	// Verify version headers are set
+	require.Equal(
+		t,
+		[]string{SDKVersion},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientVersionHeaderName),
+	)
+	require.Equal(
+		t,
+		[]string{clientNameHeaderValue},
+		metadata.ValueFromIncomingContext(srv.getSystemInfoRequestContext, clientNameHeaderName),
+	)
 }
 
 func TestCredentialsMTLS(t *testing.T) {

diff --git a/internal/internal_workflow_client.go b/internal/internal_workflow_client.go
@@ -1229,9 +1229,9 @@ func (wc *WorkflowClient) loadCapabilities(ctx context.Context, getSystemInfoTim
 	if getSystemInfoTimeout == 0 {
 		getSystemInfoTimeout = defaultGetSystemInfoTimeout
 	}
-	ctx, cancel := context.WithTimeout(ctx, getSystemInfoTimeout)
+	grpcCtx, cancel := newGRPCContext(ctx, grpcTimeout(getSystemInfoTimeout))
 	defer cancel()
-	resp, err := wc.workflowService.GetSystemInfo(ctx, &workflowservice.GetSystemInfoRequest{})
+	resp, err := wc.workflowService.GetSystemInfo(grpcCtx, &workflowservice.GetSystemInfoRequest{})
 	// We ignore unimplemented
 	if _, isUnimplemented := err.(*serviceerror.Unimplemented); err != nil && !isUnimplemented {
 		return nil, fmt.Errorf("failed reaching server: %w", err)