[FnApi] Add support header metadata in ApiServiceDescriptor. Sdks indicate they support this via StandardProtocols.API_SERVICE_DESCRIPTOR_NAME. Runners may use such metadata to distinguish between incoming fnapi streams. A concrete usage for this is to indicate for different bundles to utilize different data streams to provide isolation.

Author: scwhittle

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto

@@ -1689,6 +1689,12 @@ message StandardProtocols {
     // Indicates whether the SDK supports multimap state.
     MULTIMAP_STATE = 12
     [(beam_urn) = "beam:protocol:multimap_state:v1"];
+
+    // Indicates that the SDK supports the field ApiServiceDescriptor.header_metadata being
+    // non-empty. The header metadata specified will be added to when connecting to
+    // the service.
+    API_SERVICE_DESCRIPTOR_HEADER_METADATA = 13
+    [(beam_urn) = "beam:protocol:api_service_descriptor_header_metadata:v1"];
   }
 }
 

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/endpoints.proto

@@ -37,6 +37,11 @@ message ApiServiceDescriptor {
   // url is already being performed in a trusted context (e.g. localhost,
   // private network).
   AuthenticationSpec authentication = 2;
+
+  // (Optional) Additional ASCII header metadata that will be sent with calls to
+  // the endpoint. SDKs indicate that they handle
+  // this field with the `StandardProtocols.API_SERVICE_DESCRIPTOR_HEADER_METADATA`.
+  map<string, string> header_metadata = 3;
 }
 
 message AuthenticationSpec {

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactory.java

@@ -24,8 +24,10 @@
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientInterceptor;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Metadata;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessChannelBuilder;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.netty.NettyChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.MetadataUtils;
 import org.apache.beam.vendor.grpc.v1p69p0.io.netty.channel.epoll.EpollDomainSocketChannel;
 import org.apache.beam.vendor.grpc.v1p69p0.io.netty.channel.epoll.EpollEventLoopGroup;
 import org.apache.beam.vendor.grpc.v1p69p0.io.netty.channel.epoll.EpollSocketChannel;
@@ -95,6 +97,14 @@ public ManagedChannel forDescriptor(ApiServiceDescriptor apiServiceDescriptor) {
             // https://github.com/grpc/proposal/blob/master/A6-client-retries.md#when-retries-are-valid
             .disableRetry()
             .intercept(interceptors);
+    if (!apiServiceDescriptor.getHeaderMetadataMap().isEmpty()) {
+      Metadata metadata = new Metadata();
+      apiServiceDescriptor
+          .getHeaderMetadataMap()
+          .forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v));
+      channelBuilder =
+          channelBuilder.intercept(MetadataUtils.newAttachHeadersInterceptor(metadata));
+    }
     if (directExecutor) {
       channelBuilder = channelBuilder.directExecutor();
     }

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java

@@ -522,6 +522,8 @@ public static Set<String> getJavaCapabilities() {
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.SDK_CONSUMING_RECEIVED_DATA));
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.ORDERED_LIST_STATE));
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.MULTIMAP_STATE));
+    capabilities.add(
+        BeamUrns.getUrn(StandardProtocols.Enum.API_SERVICE_DESCRIPTOR_HEADER_METADATA));
     return capabilities.build();
   }
 

sdks/java/core/src/test/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactoryTest.java

@@ -22,7 +22,11 @@
 import static org.junit.Assume.assumeTrue;
 
 import org.apache.beam.model.pipeline.v1.Endpoints;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.CallOptions;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ClientCall;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Metadata;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.TestMethodDescriptors;
 import org.apache.commons.lang3.SystemUtils;
 import org.junit.Rule;
 import org.junit.Test;
@@ -42,6 +46,60 @@ public void testDefaultChannel() {
     ManagedChannel channel =
         ManagedChannelFactory.createDefault().forDescriptor(apiServiceDescriptor);
     assertEquals("localhost:123", channel.authority());
+    ClientCall<Void, Void> call =
+        channel.newCall(TestMethodDescriptors.voidMethod(), CallOptions.DEFAULT);
+    Metadata headers = new Metadata();
+    call.start(new ClientCall.Listener<>() {}, headers);
+    for (String key : headers.keys()) {
+      assertTrue(
+          "Unexpected key "
+              + key
+              + " with value "
+              + headers.get(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)),
+          key.startsWith("grpc-"));
+    }
+    call.cancel("Done", null);
+    channel.shutdownNow();
+  }
+
+  @Test
+  public void testDefaultChannelWithStreamName() {
+    Endpoints.ApiServiceDescriptor apiServiceDescriptor =
+        Endpoints.ApiServiceDescriptor.newBuilder()
+            .setUrl("localhost:123")
+            .putHeaderMetadata("key", "value")
+            .putHeaderMetadata("other-key", "other-value")
+            .build();
+    ManagedChannel channel =
+        ManagedChannelFactory.createDefault().forDescriptor(apiServiceDescriptor);
+    assertEquals("localhost:123", channel.authority());
+
+    ClientCall<Void, Void> call =
+        channel.newCall(TestMethodDescriptors.voidMethod(), CallOptions.DEFAULT);
+    Metadata headers = new Metadata();
+    call.start(new ClientCall.Listener<>() {}, headers);
+    boolean foundKey = false;
+    boolean foundOtherKey = false;
+    for (String key : headers.keys()) {
+      if (key.equals("key")) {
+        foundKey = true;
+        assertEquals("value", headers.get(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)));
+      } else if (key.equals("other-key")) {
+        foundOtherKey = true;
+        assertEquals(
+            "other-value", headers.get(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)));
+      } else {
+        assertTrue(
+            "Unexpected key "
+                + key
+                + " with value "
+                + headers.get(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)),
+            key.startsWith("grpc-"));
+      }
+    }
+    assertTrue(foundKey);
+    assertTrue(foundOtherKey);
+    call.cancel("Done", null);
     channel.shutdownNow();
   }
 

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java

@@ -110,6 +110,6 @@ private BeamFnDataGrpcMultiplexer getClientFor(
             new BeamFnDataGrpcMultiplexer(
                 descriptor,
                 outboundObserverFactory,
-                BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor))::data));
+                BeamFnDataGrpc.newStub(channelFactory.apply(descriptor))::data));
   }
 }