Add reverse propagation of telemetry tags

src/main/azure-functions-language-worker-protobuf/src/proto/FunctionRpc.proto

@@ -245,6 +245,14 @@ message FunctionEnvironmentReloadRequest {
 }
 
 message FunctionEnvironmentReloadResponse {
+  enum CapabilitiesUpdateStrategy {
+    // overwrites existing values and appends new ones
+    // ex. worker init: {A: foo, B: bar} + env reload: {A:foo, B: foo, C: foo} -> {A: foo, B: foo, C: foo}
+    merge = 0;
+    // existing capabilities are cleared and new capabilities are applied
+    // ex. worker init: {A: foo, B: bar} + env reload: {A:foo, C: foo} -> {A: foo, C: foo}
+    replace = 1;
+  }
   // After specialization, worker sends capabilities & metadata.
   // Worker metadata captured for telemetry purposes
   WorkerMetadata worker_metadata = 1;
@@ -254,6 +262,9 @@ message FunctionEnvironmentReloadResponse {
 
   // Status of the response
   StatusResult result = 3;
+
+  // If no strategy is defined, the host will default to merge
+  CapabilitiesUpdateStrategy capabilities_update_strategy = 4;
 }
 
 // Tell the out-of-proc worker to close any shared memory maps it allocated for given invocation
@@ -337,6 +348,9 @@ message RpcFunctionMetadata {
   // A flag indicating if managed dependency is enabled or not
   bool managed_dependency_enabled = 14;
 
+  // The optional function execution retry strategy to use on invocation failures.
+  RpcRetryOptions retry_options = 15;
+
   // Properties for function metadata
   // They're usually specific to a worker and largely passed along to the controller API for use
   // outside the host
@@ -423,11 +437,15 @@ message InvocationResponse {
   // Output binding data
   repeated ParameterBinding output_data = 2;
 
+  // Status of the invocation (success/failure/canceled)
+  StatusResult result = 3;
+
   // data returned from Function (for $return and triggers with return support)
   TypedData return_value = 4;
 
-  // Status of the invocation (success/failure/canceled)
-  StatusResult result = 3;
+  // Enables propagation of tags from worker -> host.
+  map<string, string> trace_context_attributes = 5;
+
 }
 
 message WorkerWarmupRequest {
@@ -466,8 +484,8 @@ enum RpcDataType {
   bytes = 3;
   stream = 4;
   http = 5;
-  int_x = 6;
-  double_x = 7;
+  int = 6;
+  double = 7;
   collection_bytes = 8;
   collection_string = 9;
   collection_double = 10;
@@ -698,3 +716,30 @@ message ModelBindingData
 message CollectionModelBindingData {
   repeated ModelBindingData model_binding_data = 1;
 }
+
+// Retry policy which the worker sends the host when the worker indexes
+// a function.
+message RpcRetryOptions
+{
+  // The retry strategy to use. Valid values are fixed delay or exponential backoff.
+  enum RetryStrategy
+  {
+      exponential_backoff = 0;
+      fixed_delay = 1;
+  }
+
+  // The maximum number of retries allowed per function execution.
+  //  -1 means to retry indefinitely.
+  int32 max_retry_count = 2;
+
+  // The delay that's used between retries when you're using a fixed delay strategy.
+  google.protobuf.Duration delay_interval = 3;
+
+  // The minimum retry delay when you're using an exponential backoff strategy
+  google.protobuf.Duration minimum_interval = 4;
+
+  // The maximum retry delay when you're using an exponential backoff strategy
+  google.protobuf.Duration maximum_interval = 5;
+
+  RetryStrategy retry_strategy = 6;
+}

src/main/azure-functions-language-worker-protobuf/src/proto/servicebus/settlement.proto

@@ -0,0 +1,99 @@
+syntax = "proto3";
+
+import "google/protobuf/empty.proto";
+import "google/protobuf/wrappers.proto";
+import "google/protobuf/timestamp.proto";
+
+// this namespace will be shared between isolated worker and WebJobs extension so make it somewhat generic
+option csharp_namespace = "Microsoft.Azure.ServiceBus.Grpc";
+
+// The settlement service definition.
+service Settlement {
+  // Completes a message
+  rpc Complete (CompleteRequest) returns (google.protobuf.Empty) {}
+
+  // Abandons a message
+  rpc Abandon (AbandonRequest) returns (google.protobuf.Empty) {}
+
+  // Deadletters a message
+  rpc Deadletter (DeadletterRequest) returns (google.protobuf.Empty) {}
+
+  // Defers a message
+  rpc Defer (DeferRequest) returns (google.protobuf.Empty) {}
+
+  // Renew message lock
+  rpc RenewMessageLock (RenewMessageLockRequest) returns (google.protobuf.Empty) {}
+
+  // Get session state
+  rpc GetSessionState (GetSessionStateRequest) returns (GetSessionStateResponse) {}
+
+  // Set session state
+  rpc SetSessionState (SetSessionStateRequest) returns (google.protobuf.Empty) {}
+
+  // Release session
+  rpc ReleaseSession (ReleaseSessionRequest) returns (google.protobuf.Empty) {}
+
+  // Renew session lock
+  rpc RenewSessionLock (RenewSessionLockRequest) returns (RenewSessionLockResponse) {}
+}
+
+// The complete message request containing the locktoken.
+message CompleteRequest {
+  string locktoken = 1;
+}
+
+// The abandon message request containing the locktoken and properties to modify.
+message AbandonRequest {
+  string locktoken = 1;
+  bytes propertiesToModify = 2;
+}
+
+// The deadletter message request containing the locktoken and properties to modify along with the reason/description.
+message DeadletterRequest {
+  string locktoken = 1;
+  bytes propertiesToModify = 2;
+  google.protobuf.StringValue deadletterReason = 3;
+  google.protobuf.StringValue deadletterErrorDescription = 4;
+}
+
+// The defer message request containing the locktoken and properties to modify.
+message DeferRequest {
+  string locktoken = 1;
+  bytes propertiesToModify = 2;
+}
+
+// The renew message lock request containing the locktoken.
+message RenewMessageLockRequest {
+  string locktoken = 1;
+}
+
+// The get message request.
+message GetSessionStateRequest {
+  string sessionId = 1;
+}
+
+// The set message request.
+message SetSessionStateRequest {
+  string sessionId = 1;
+  bytes sessionState = 2;
+}
+
+// Get response containing the session state.
+message GetSessionStateResponse {
+  bytes sessionState = 1;
+}
+
+// Release session.
+message ReleaseSessionRequest {
+  string sessionId = 1;
+}
+
+// Renew session lock.
+message RenewSessionLockRequest {
+  string sessionId = 1;
+}
+
+// Renew session lock.
+message RenewSessionLockResponse {
+  google.protobuf.Timestamp lockedUntil = 1;
+}

src/main/java/com/microsoft/azure/functions/worker/broker/JavaFunctionBroker.java

@@ -187,9 +187,10 @@ private FunctionExecutionMiddleware getFunctionExecutionMiddleWare(ClassLoader c
 		return functionExecutionMiddleware;
 	}
 
-	public Optional<TypedData> invokeMethod(String id, InvocationRequest request, List<ParameterBinding> outputs)
+	public Optional<TypedData> invokeMethod(String id, InvocationRequest request, List<ParameterBinding> outputs,
+			Map<String, String> traceContextAttributes)
 			throws Exception {
-		ExecutionContextDataSource executionContextDataSource = buildExecutionContext(id, request);
+		ExecutionContextDataSource executionContextDataSource = buildExecutionContext(id, request, traceContextAttributes);
 
 		if (isJavaSdkTypesEnabled()) {
 			this.functionFactories.get(id).create().doNext(executionContextDataSource);
@@ -201,7 +202,8 @@ public Optional<TypedData> invokeMethod(String id, InvocationRequest request, Li
 		return executionContextDataSource.getDataStore().getDataTargetTypedValue(BindingDataStore.RETURN_NAME);
 	}
 
-	private ExecutionContextDataSource buildExecutionContext(String id,  InvocationRequest request)
+	private ExecutionContextDataSource buildExecutionContext(String id, InvocationRequest request,
+			Map<String, String> traceContextAttributes)
 			throws NoSuchMethodException {
 		ImmutablePair<String, FunctionDefinition> methodEntry = this.methods.get(id);
 		FunctionDefinition functionDefinition = methodEntry.right;
@@ -213,7 +215,7 @@ private ExecutionContextDataSource buildExecutionContext(String id,  InvocationR
 		dataStore.addTriggerMetadataSource(getTriggerMetadataMap(request));
 		dataStore.addParameterSources(request.getInputDataList());
 		ExecutionTraceContext traceContext = new ExecutionTraceContext(request.getTraceContext().getTraceParent(),
-				request.getTraceContext().getTraceState(), request.getTraceContext().getAttributesMap());
+				request.getTraceContext().getTraceState(), traceContextAttributes);
 		ExecutionRetryContext retryContext = new ExecutionRetryContext(request.getRetryContext().getRetryCount(),
 				request.getRetryContext().getMaxRetryCount(), request.getRetryContext().getException());
 		ExecutionContextDataSource executionContextDataSource = new ExecutionContextDataSource(

src/main/java/com/microsoft/azure/functions/worker/handler/InvocationRequestHandler.java

@@ -31,8 +31,10 @@ String execute(InvocationRequest request, InvocationResponse.Builder response) t
         response.setInvocationId(invocationId);
 
         List<ParameterBinding> outputBindings = new ArrayList<>();
-        this.broker.invokeMethod(functionId, request, outputBindings).ifPresent(response::setReturnValue);
+        Map<String, String> traceContextAttributes = new HashMap<>(request.getTraceContext().getAttributesMap());
+        this.broker.invokeMethod(functionId, request, outputBindings, traceContextAttributes).ifPresent(response::setReturnValue);
         response.addAllOutputData(outputBindings);
+        response.putAllTraceContextAttributes(traceContextAttributes);
 
         return String.format("Function \"%s\" (Id: %s) invoked by Java Worker",
                 this.broker.getMethodName(functionId).orElse("UNKNOWN"), invocationId);

src/test/java/com/microsoft/azure/functions/worker/functional/tests/SimpleParamReturnTest.java

@@ -1,5 +1,6 @@
 package com.microsoft.azure.functions.worker.functional.tests;
 
+import com.microsoft.azure.functions.ExecutionContext;
 import com.microsoft.azure.functions.rpc.messages.*;
 import com.microsoft.azure.functions.worker.test.utilities.*;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -14,6 +15,11 @@ public String ReturnStringFunction() {
 		return stringReturnValue;
 	}
 
+	public String AddTraceContextAttributeFunction(ExecutionContext context) {
+		context.getTraceContext().getAttributes().put("customKey", "customValue");
+		return stringReturnValue;
+	}
+
 	@ParameterizedTest
 	@ValueSource(strings = {"test String"})
 	public void testStringData(String stringInput) throws Exception {
@@ -26,4 +32,16 @@ public void testStringData(String stringInput) throws Exception {
 			assertEquals(stringInput, stringResponse.getReturnValue().getString());
 		}
 	}
+
+	@ParameterizedTest
+	@ValueSource(strings = {"test String"})
+	public void testTraceContextAttributesAreAddedToInvocationResponse(String stringInput) throws Exception {
+		stringReturnValue = stringInput;
+		System.setProperty("azure.functions.worker.java.skip.testing", "true");
+		try (FunctionsTestHost host = new FunctionsTestHost()) {
+			this.loadFunction(host, "traceContextTestId", "AddTraceContextAttributeFunction");
+			InvocationResponse response = host.call("traceContext", "traceContextTestId");
+			assertEquals("customValue", response.getTraceContextAttributesMap().get("customKey"));
+		}
+	}
 }