You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "m-trieu (via GitHub)" <gi...@apache.org> on 2023/08/01 03:59:02 UTC

[GitHub] [beam] m-trieu opened a new pull request, #27767: add getworkermetadata streaming rpc

m-trieu opened a new pull request, #27767:
URL: https://github.com/apache/beam/pull/27767

   Adding a stream for GetWorkerMetadata
   
   Does not include (the following will be in a seperate CL)
   1. Storing the metadata
   2. Using the metadata
   3. Adding stream to GrpcWindmillServer
   
   - Updated the type parameters for the *ResponseObserver.java files
   - Extend AbstractWindmillStream, but modify some behavior to support client unary request/server stream response
   - Add unit tests (should do this for all of the stream classes)
   
   r: @scwhittle @dustin12 
   
   ------------------------
   
   Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
   
    - [ ] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
    - [ ] Update `CHANGES.md` with noteworthy changes.
    - [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
   
   See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/get-started-contributing/#make-the-reviewers-job-easier).
   
   To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
   
   GitHub Actions Tests Status (on master branch)
   ------------------------------------------------------------------------------------------------
   [![Build python source distribution and wheels](https://github.com/apache/beam/workflows/Build%20python%20source%20distribution%20and%20wheels/badge.svg?branch=master&event=schedule)](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
   [![Python tests](https://github.com/apache/beam/workflows/Python%20tests/badge.svg?branch=master&event=schedule)](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Java tests](https://github.com/apache/beam/workflows/Java%20Tests/badge.svg?branch=master&event=schedule)](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Go tests](https://github.com/apache/beam/workflows/Go%20tests/badge.svg?branch=master&event=schedule)](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
   
   See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1313245206


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+
+  private static final Logger LOG = LoggerFactory.getLogger(WindmillEndpoints.class);
+
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        parseDirectEndpoint(endpointProto.getDirectEndpoint())
+            .ifPresent(endpointBuilder::setDirectEndpoint);
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+
+      Endpoint endpoint = endpointBuilder.build();
+
+      if (!endpoint.directEndpoint().isPresent() && !endpoint.workerToken().isPresent()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "direct_endpoint=[%s] not present or could not be parsed, and worker_token"
+                    + " not present. At least one of these fields is required.",
+                endpointProto.getDirectEndpoint()));
+      }
+
+      return endpoint;
+    }
+
+    private static Optional<Inet6Address> parseDirectEndpoint(String directEndpoint) {
+      InetAddress directEndpointAddress = null;
+      try {
+        directEndpointAddress = Inet6Address.getByName(directEndpoint);
+      } catch (UnknownHostException e) {
+        LOG.warn(
+            "Error occurred trying to parse direct_endpoint={} into IPv6 address. Exception={}",
+            directEndpoint,
+            e.toString());
+      }
+
+      // Inet6Address.getByAddress returns either an IPv4 or an IPv6 address depending on the format
+      // of the direct_endpoint string.
+      if (!(directEndpointAddress instanceof Inet6Address)) {
+        LOG.warn(
+            "{} is not an IPv6 address. Direct endpoints are expected to be in IPv6 format.",
+            directEndpoint);
+      }
+
+      return directEndpointAddress == null
+          ? Optional.empty()
+          : Optional.of((Inet6Address) directEndpointAddress);

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1312770471


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));
+    } else {
+      // If the currentMetadataVersion is greater than or equal to one in the response, the response
+      // data is stale, and we do not want to do anything.
+      LOG.info(
+          "Received metadata version={}; Current metadata version={}. "
+              + "Skipping update because received stale metadata",
+          response.getMetadataVersion(),
+          metadataVersion);
+    }
+
+    this.latestResponse = response;

Review Comment:
   don't update lastResponse in the case it is stale



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   Looking again, the AbstractWindmillStream are used for longer than the deadline, they reconnect grpc streams internally.  So I think you need to use the observer approach and not a fixed deadline.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+
+  private static final Logger LOG = LoggerFactory.getLogger(WindmillEndpoints.class);
+
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        parseDirectEndpoint(endpointProto.getDirectEndpoint())
+            .ifPresent(endpointBuilder::setDirectEndpoint);
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+
+      Endpoint endpoint = endpointBuilder.build();
+
+      if (!endpoint.directEndpoint().isPresent() && !endpoint.workerToken().isPresent()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "direct_endpoint=[%s] not present or could not be parsed, and worker_token"
+                    + " not present. At least one of these fields is required.",
+                endpointProto.getDirectEndpoint()));
+      }
+
+      return endpoint;
+    }
+
+    private static Optional<Inet6Address> parseDirectEndpoint(String directEndpoint) {
+      InetAddress directEndpointAddress = null;
+      try {
+        directEndpointAddress = Inet6Address.getByName(directEndpoint);
+      } catch (UnknownHostException e) {
+        LOG.warn(
+            "Error occurred trying to parse direct_endpoint={} into IPv6 address. Exception={}",
+            directEndpoint,
+            e.toString());
+      }
+
+      // Inet6Address.getByAddress returns either an IPv4 or an IPv6 address depending on the format
+      // of the direct_endpoint string.
+      if (!(directEndpointAddress instanceof Inet6Address)) {
+        LOG.warn(
+            "{} is not an IPv6 address. Direct endpoints are expected to be in IPv6 format.",
+            directEndpoint);
+      }
+
+      return directEndpointAddress == null
+          ? Optional.empty()
+          : Optional.of((Inet6Address) directEndpointAddress);

Review Comment:
   this cast seems wrong if the above instanceof fails.
   seems we should return empty in that case
   
   since null will be false for instanceof you could just do a single check using that.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));
+    } else {
+      // If the currentMetadataVersion is greater than or equal to one in the response, the response
+      // data is stale, and we do not want to do anything.
+      LOG.info(
+          "Received metadata version={}; Current metadata version={}. "
+              + "Skipping update because received stale metadata",
+          response.getMetadataVersion(),
+          metadataVersion);
+    }
+
+    this.latestResponse = response;
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected boolean hasPendingRequests() {
+    return false;
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void sendHealthCheck() {
+    send(HEALTH_CHECK_REQUEST);
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s], latest_response=[%s]",
+        metadataVersion, workerMetadataRequest.getHeader(), latestResponse);

Review Comment:
   I think you need some synchronization for latestResponse etc since I believe the html rendering can be concurrent



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));

Review Comment:
   if you add sychronization for html rendering (see below), it would be good to keep the accept call out of the synchronized block.



##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,18 +758,15 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    // IPv6 address of a streaming engine windmill worker.
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;
   }
   repeated Endpoint work_endpoints = 2;
 
   // Maps from GlobalData tag to the endpoint that should be used for GetData
   // calls to retrieve that global data.
   map<string, Endpoint> global_data_endpoints = 3;
-
-  // DirectPath endpoints to be used by user workers for streaming engine jobs.
-  // DirectPath endpoints here are virtual IPv6 addresses of the windmill
-  // workers.
-  repeated Endpoint direct_path_endpoints = 4;

Review Comment:
   reserve the tag # so it is not reused
   
   reserved 4;



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));

Review Comment:
   From AbstractStream header it sounds like you should not synchronize on this but some other object:
   
    * <p>Synchronization on this is used to synchronize the gRpc stream state and internal data
    * structures. Since grpc channel operations may block, synchronization on this stream may also
    * block. This is generally not a problem since streams are used in a single-threaded manner.
    * However, some accessors used for status page and other debugging need to take care not to require
    * synchronizing on this.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1291187557


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -133,13 +130,13 @@ private static long debugDuration(long nowMs, long startMs) {
   protected abstract boolean hasPendingRequests();
 
   /**
-   * Called when the stream is throttled due to resource exhausted errors. Will be called for each
-   * resource exhausted error not just the first. onResponse() must stop throttling on receipt of
-   * the first good message.
+   * Called when the client side stream is throttled due to resource exhausted errors. Will be
+   * called for each resource exhausted error not just the first. onResponse() must stop throttling
+   * on receipt of the first good message.
    */
   protected abstract void startThrottleTimer();
 
-  private StreamObserver<RequestT> requestObserver() {
+  protected StreamObserver<RequestT> requestObserver() {

Review Comment:
   can this remain private? don't see new call site



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -252,7 +249,7 @@ public final Instant startTime() {
     return new Instant(startTimeMs.get());
   }
 
-  private class ResponseObserver implements StreamObserver<ResponseT> {
+  public class ResponseObserver implements StreamObserver<ResponseT> {

Review Comment:
   ditto, seems like this could remain private?



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -149,7 +146,7 @@ private StreamObserver<RequestT> requestObserver() {
   }
 
   /** Send a request to the server. */
-  protected final void send(RequestT request) {
+  protected void send(RequestT request) {

Review Comment:
   can this remain final? don't see new override



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1679376719

   changes to add worker token support are independent, this review can proceed 
   thanks! @scwhittle 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1316558104


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));
+    } else {
+      // If the currentMetadataVersion is greater than or equal to one in the response, the response
+      // data is stale, and we do not want to do anything.
+      LOG.info(
+          "Received metadata version={}; Current metadata version={}. "
+              + "Skipping update because received stale metadata",
+          response.getMetadataVersion(),
+          metadataVersion);
+    }
+
+    this.latestResponse = response;
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected boolean hasPendingRequests() {
+    return false;
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void sendHealthCheck() {
+    send(HEALTH_CHECK_REQUEST);
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s], latest_response=[%s]",
+        metadataVersion, workerMetadataRequest.getHeader(), latestResponse);

Review Comment:
   done. added locking



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1319045632


##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,336 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final String IPV6_ADDRESS_1 = "2001:db8:0000:bac5:0000:0000:fed0:81a2";
+  private static final String IPV6_ADDRESS_2 = "2001:db8:0000:bac5:0000:0000:fed0:82a3";
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder()
+              .setDirectEndpoint(IPV6_ADDRESS_1)
+              .setWorkerToken("worker_token")
+              .build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      int metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.create(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        FluentBackoff.DEFAULT.backoff(),
+        StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+        streamRegistry,
+        1, // logEveryNStreamFailures
+        TEST_JOB_HEADER,
+        metadataVersion,
+        new ThrottleTimer(),
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder()
+            .setDirectEndpoint(IPV6_ADDRESS_1)
+            .setWorkerToken("worker_token")
+            .build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllWorkEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    GetWorkerMetadataTestStub testStub =
+        new GetWorkerMetadataTestStub(new TestGetWorkMetadataRequestObserver());
+    int metadataVersion = -1;
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    testStub.injectWorkerMetadata(mockResponse);
+
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(

Review Comment:
   are there matchers you could use instead? This is where having someone more familiar with Java could help review.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308261178


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   i rebased this branch onto the branch that removes the sync stubs so i could reuse some of the code.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1283922827


##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,335 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String name = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(name)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(InProcessChannelBuilder.forName(name).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(GLOBAL_DATA_ENDPOINTS.containsKey(key)));
+    assertEquals(
+        DIRECT_PATH_ENDPOINTS.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            DIRECT_PATH_ENDPOINTS.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(mockResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_consumesSubsequentResponseMetadata() {
+    WorkerMetadataResponse initialResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(initialResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+
+    List<WorkerMetadataResponse.Endpoint> newDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("newWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> newGlobalDataEndpoints = new HashMap<>();
+    newGlobalDataEndpoints.put(
+        "new_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("newGlobalData").build());
+
+    WorkerMetadataResponse newWorkMetadataResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(initialResponse.getMetadataVersion() + 1)
+            .addAllDirectPathEndpoints(newDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(newGlobalDataEndpoints)
+            .build();
+
+    testStub.injectWorkerMetadata(newWorkMetadataResponse);
+
+    assertEquals(
+        newGlobalDataEndpoints.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(newGlobalDataEndpoints.containsKey(key)));
+    assertEquals(
+        newDirectPathEndpoints.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            newDirectPathEndpoints.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(newWorkMetadataResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_doesNotConsumeResponseIfMetadataStale() {
+    WorkerMetadataResponse freshEndpoints =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(2)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(freshEndpoints);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    List<WorkerMetadataResponse.Endpoint> staleDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("staleWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> staleGlobalDataEndpoints = new HashMap<>();
+    staleGlobalDataEndpoints.put(
+        "stale_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("staleGlobalData").build());
+
+    testStub.injectWorkerMetadata(
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(staleDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(staleGlobalDataEndpoints)
+            .build());
+
+    // Should have ignored the stale update and only used initial.
+    verify(testWindmillEndpointsConsumer)
+        .accept(
+            WindmillEndpoints.builder()
+                .setWindmillServers(
+                    freshEndpoints.getDirectPathEndpointsList().stream()
+                        .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                        .collect(toImmutableList()))
+                .setGlobalDataServers(
+                    freshEndpoints.getGlobalDataEndpointsMap().entrySet().stream()
+                        .collect(
+                            toImmutableMap(Map.Entry::getKey, e -> e.getValue().getEndpoint())))
+                .build());
+    verifyNoMoreInteractions(testWindmillEndpointsConsumer);
+    assertEquals(freshEndpoints.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_correctlyAddsAndRemovesStreamFromRegistry() {
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(
+            WorkerMetadataResponse.newBuilder()
+                .setMetadataVersion(1)
+                .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+                .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+                .build());
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    stream =
+        getWorkerMetadataStream(testStub, new AtomicLong(), new TestWindmillEndpointsConsumer());
+
+    assertTrue(streamRegistry.contains(stream));
+    stream.close();
+    assertFalse(streamRegistry.contains(stream));
+  }
+
+  @Test
+  public void testSendHealthCheck() {

Review Comment:
   will add behavior here once we decide



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1291778190


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -252,7 +249,7 @@ public final Instant startTime() {
     return new Instant(startTimeMs.get());
   }
 
-  private class ResponseObserver implements StreamObserver<ResponseT> {
+  public class ResponseObserver implements StreamObserver<ResponseT> {

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290771660


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,189 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        getWorkerMetadataRpc,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private long updateMetadataVersion(long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    metadataVersion.getAndUpdate(
+        current -> updateMetadataVersion(current, response.getMetadataVersion()));

Review Comment:
   done. good catch! tried to use the Atomic* non blocking to do the filtering.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290773607


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcCommitWorkStream.java:
##########
@@ -56,7 +56,8 @@ final class GrpcCommitWorkStream
   private final int streamingRpcBatchLimit;
 
   private GrpcCommitWorkStream(
-      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      Function<StreamObserver<StreamingCommitResponse>, StreamObserver<StreamingCommitWorkRequest>>
+          commitWorkRpc,

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1313392298


##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,18 +758,15 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    // IPv6 address of a streaming engine windmill worker.
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;
   }
   repeated Endpoint work_endpoints = 2;
 
   // Maps from GlobalData tag to the endpoint that should be used for GetData
   // calls to retrieve that global data.
   map<string, Endpoint> global_data_endpoints = 3;
-
-  // DirectPath endpoints to be used by user workers for streaming engine jobs.
-  // DirectPath endpoints here are virtual IPv6 addresses of the windmill
-  // workers.
-  repeated Endpoint direct_path_endpoints = 4;

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1316558757


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));

Review Comment:
   i opted to use StampedLock https://docs.oracle.com/javase/8/docs/api/java/util/concurrent/locks/StampedLock.html cause it supports optimistic locking and allows for multiple readers.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1317062950


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+
+  /** Guards access to metadataVersion and latestResponse. */
+  private final StampedLock metadataLock;
+
+  private long metadataVersion;

Review Comment:
   use GuardedBy annotations



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+
+  /** Guards access to metadataVersion and latestResponse. */
+  private final StampedLock metadataLock;
+
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer,
+      StampedLock metadataLock) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+    this.metadataLock = metadataLock;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater,
+            new StampedLock());
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater,
+      StampedLock lock) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater,
+            lock);

Review Comment:
   what about just using @VisibleForTesting annotation instead of injecting lower-level things like the lock which seems like it adds boilerplate here and in the test as well as making it less obvious that it is just for testing.
   
   or add a getMetadataVersion() method that the tests use



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+
+  /** Guards access to metadataVersion and latestResponse. */
+  private final StampedLock metadataLock;
+
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer,
+      StampedLock metadataLock) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+    this.metadataLock = metadataLock;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater,
+            new StampedLock());
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater,
+      StampedLock lock) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater,
+            lock);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse().
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long readStamp = metadataLock.tryOptimisticRead();

Review Comment:
   I don't think the complication of the StampedLock is beneficial here. Both the responses and status page are fairly infrequent and the logic in the critical section is minimal.
   
   More complex locks can perform worse for such cases and is certainly more error-prone.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290774775


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -294,7 +293,10 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.size() == 1) {
       return stubList.get(0);
     }
-    return stubList.get(rand.nextInt(stubList.size()));
+    return stubList
+        .get(rand.nextInt(stubList.size()))
+        .withDeadlineAfter(
+            AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1291779745


##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,348 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String name = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(name)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(InProcessChannelBuilder.forName(name).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(GLOBAL_DATA_ENDPOINTS.containsKey(key)));
+    assertEquals(
+        DIRECT_PATH_ENDPOINTS.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            DIRECT_PATH_ENDPOINTS.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(mockResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_consumesSubsequentResponseMetadata() {
+    WorkerMetadataResponse initialResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(initialResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+
+    List<WorkerMetadataResponse.Endpoint> newDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("newWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> newGlobalDataEndpoints = new HashMap<>();
+    newGlobalDataEndpoints.put(
+        "new_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("newGlobalData").build());
+
+    WorkerMetadataResponse newWorkMetadataResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(initialResponse.getMetadataVersion() + 1)
+            .addAllDirectPathEndpoints(newDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(newGlobalDataEndpoints)
+            .build();
+
+    testStub.injectWorkerMetadata(newWorkMetadataResponse);
+
+    assertEquals(
+        newGlobalDataEndpoints.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(newGlobalDataEndpoints.containsKey(key)));
+    assertEquals(
+        newDirectPathEndpoints.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            newDirectPathEndpoints.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(newWorkMetadataResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_doesNotConsumeResponseIfMetadataStale() {
+    WorkerMetadataResponse freshEndpoints =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(2)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(freshEndpoints);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    List<WorkerMetadataResponse.Endpoint> staleDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("staleWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> staleGlobalDataEndpoints = new HashMap<>();
+    staleGlobalDataEndpoints.put(
+        "stale_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("staleGlobalData").build());
+
+    testStub.injectWorkerMetadata(

Review Comment:
   it failed (i didnt rerun it after i removed the locking)
   passes now that the logic is fixed



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1291777891


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -149,7 +146,7 @@ private StreamObserver<RequestT> requestObserver() {
   }
 
   /** Send a request to the server. */
-  protected final void send(RequestT request) {
+  protected void send(RequestT request) {

Review Comment:
   missed it in the revert done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] github-actions[bot] commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1674270979

   Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:
   
   R: @lostluck added as fallback since no labels match configuration
   
   Available commands:
   - `stop reviewer notifications` - opt out of the automated review tooling
   - `remind me after tests pass` - tag the comment author after tests pass
   - `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
   
   The PR bot will only process comments in the main thread (not review comments).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1319228520


##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,336 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final String IPV6_ADDRESS_1 = "2001:db8:0000:bac5:0000:0000:fed0:81a2";
+  private static final String IPV6_ADDRESS_2 = "2001:db8:0000:bac5:0000:0000:fed0:82a3";
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder()
+              .setDirectEndpoint(IPV6_ADDRESS_1)
+              .setWorkerToken("worker_token")
+              .build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      int metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.create(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        FluentBackoff.DEFAULT.backoff(),
+        StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+        streamRegistry,
+        1, // logEveryNStreamFailures
+        TEST_JOB_HEADER,
+        metadataVersion,
+        new ThrottleTimer(),
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder()
+            .setDirectEndpoint(IPV6_ADDRESS_1)
+            .setWorkerToken("worker_token")
+            .build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllWorkEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    GetWorkerMetadataTestStub testStub =
+        new GetWorkerMetadataTestStub(new TestGetWorkMetadataRequestObserver());
+    int metadataVersion = -1;
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    testStub.injectWorkerMetadata(mockResponse);
+
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(

Review Comment:
   found them adding the dependency for the gradle file.  I will update the tests to reflect and then submit
   
   using google.Truth



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308386746


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   done, decided to keep it the same and change the name for simplicity



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));
+
+    if (updatedMetadataVersion > currentMetadataVersion) {
+      serverMappingUpdater.accept(WindmillEndpoints.from(response));

Review Comment:
   done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1694029640

   @scwhittle do you have any more comments for this pull request? Thank you!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308221807


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));
+
+    if (updatedMetadataVersion > currentMetadataVersion) {
+      serverMappingUpdater.accept(WindmillEndpoints.from(response));
+    }
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected boolean hasPendingRequests() {
+    return false;
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void sendHealthCheck() {
+    send(HEALTH_CHECK_REQUEST);
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s]",

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308133634


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        endpointBuilder.setDirectEndpoint(endpointProto.getDirectEndpoint());
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+      return endpointBuilder.build();
+    }
+
+    public abstract Optional<String> directEndpoint();

Review Comment:
   lets use Inet6Address, i will update
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302358226


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,

Review Comment:
   done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302355008


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underlying stream observer that is being wrapped.
  */
-final class ForwardingClientResponseObserver<ReqT, RespT>
-    implements ClientResponseObserver<RespT, ReqT> {
+final class ForwardingClientResponseObserver<ResponseT, RequestT>

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1317710425


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+
+  /** Guards access to metadataVersion and latestResponse. */
+  private final StampedLock metadataLock;
+
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer,
+      StampedLock metadataLock) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+    this.metadataLock = metadataLock;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater,
+            new StampedLock());
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater,
+      StampedLock lock) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater,
+            lock);

Review Comment:
   done, just reused the create method in the test.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] github-actions[bot] commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1659596586

   Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1280088296


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,285 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final GetWorkMetadataStreamRequestObserver getWorkMetadataStreamRequestObserver;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        // This is a client unary request, server stream response RPC. There is no need to
+        // maintain the client side stream.
+        responseObserver -> null,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.getWorkMetadataStreamRequestObserver =
+        GetWorkMetadataStreamRequestObserver.newUnopenedStream(stub);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.fullyInitializeRequestObserver();
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.serverStream(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.fullyInitializeRequestObserver();
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Because {@link AbstractWindmillStream.ResponseObserver} is a non-static inner class, it (and
+   * its members) cannot be referenced until the constructor finishes executing and the object is
+   * completely instantiated. We gate this from being null by making the constructor private and
+   * exposing a static method so that an instance of {@link GrpcGetWorkerMetadataStream} will be
+   * created in the correct state.
+   */
+  private void fullyInitializeRequestObserver() {
+    getWorkMetadataStreamRequestObserver.init(
+        AbstractWindmillStream.ResponseObserver::new,
+        () -> {
+          synchronized (GrpcGetWorkerMetadataStream.this) {
+            streamClosed.set(true);
+          }
+        });
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> createNewRequestObserver() {
+    // we do not need to re-initiate the requestObserver here, since every
+    // GetWorkMetadataStreamRequestObserver.onNext() starts a new server stream.
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    synchronized (this) {
+      // If the metadataVersion newer or the same as the one in the response, the response data is
+      // stale, and we do not want to do anything.
+      long currentMetadataVersion = metadataVersion.get();
+      if (currentMetadataVersion >= response.getMetadataVersion()) {
+        LOG.info(
+            "Received metadata version={}; Current metadata version={}. "
+                + "Skipping update because received stale metadata",
+            response.getMetadataVersion(),
+            currentMetadataVersion);
+        return;
+      }
+      metadataVersion.set(response.getMetadataVersion());
+    }
+    ImmutableMap<String, String> globalDataServers =
+        response.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> endpoint.getValue().getEndpoint()));
+
+    ImmutableList<String> windmillServers =
+        response.getDirectPathEndpointsList().stream()
+            .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+            .collect(toImmutableList());
+
+    serverMappingUpdater.accept(
+        WindmillEndpoints.builder()
+            .setGlobalDataServers(globalDataServers)
+            .setWindmillServers(windmillServers)
+            .build());
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();

Review Comment:
   @scwhittle are these throttle timers for client side throttling? if yes, i don't think we need them here.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] lostluck merged pull request #27767: add getworkermetadata streaming rpc

Posted by "lostluck (via GitHub)" <gi...@apache.org>.
lostluck merged PR #27767:
URL: https://github.com/apache/beam/pull/27767


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1281757690


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -73,9 +73,9 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final Executor executor;
   private final BackOff backoff;
   // Indicates if the current stream in requestObserver is closed by calling close() method
-  private final AtomicBoolean streamClosed;
+  protected AtomicBoolean streamClosed;

Review Comment:
   can this be reverted?
   It seems at least it should be final, but also setting it as part of stream callback doesn't match the comment



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underling stream observer that is being wrapped.

Review Comment:
   underlying or wrapped



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -73,9 +73,9 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final Executor executor;
   private final BackOff backoff;
   // Indicates if the current stream in requestObserver is closed by calling close() method
-  private final AtomicBoolean streamClosed;
+  protected AtomicBoolean streamClosed;
   private final AtomicLong startTimeMs;
-  private final AtomicLong lastSendTimeMs;
+  protected final AtomicLong lastSendTimeMs;

Review Comment:
   can this be reverted? don't see it accessed by subclass



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,285 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final GetWorkMetadataStreamRequestObserver getWorkMetadataStreamRequestObserver;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        // This is a client unary request, server stream response RPC. There is no need to
+        // maintain the client side stream.
+        responseObserver -> null,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.getWorkMetadataStreamRequestObserver =
+        GetWorkMetadataStreamRequestObserver.newUnopenedStream(stub);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.fullyInitializeRequestObserver();
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.serverStream(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.fullyInitializeRequestObserver();
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Because {@link AbstractWindmillStream.ResponseObserver} is a non-static inner class, it (and
+   * its members) cannot be referenced until the constructor finishes executing and the object is
+   * completely instantiated. We gate this from being null by making the constructor private and
+   * exposing a static method so that an instance of {@link GrpcGetWorkerMetadataStream} will be
+   * created in the correct state.
+   */
+  private void fullyInitializeRequestObserver() {
+    getWorkMetadataStreamRequestObserver.init(
+        AbstractWindmillStream.ResponseObserver::new,
+        () -> {
+          synchronized (GrpcGetWorkerMetadataStream.this) {
+            streamClosed.set(true);
+          }
+        });
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> createNewRequestObserver() {
+    // we do not need to re-initiate the requestObserver here, since every
+    // GetWorkMetadataStreamRequestObserver.onNext() starts a new server stream.
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    synchronized (this) {
+      // If the metadataVersion newer or the same as the one in the response, the response data is
+      // stale, and we do not want to do anything.
+      long currentMetadataVersion = metadataVersion.get();
+      if (currentMetadataVersion >= response.getMetadataVersion()) {
+        LOG.info(
+            "Received metadata version={}; Current metadata version={}. "
+                + "Skipping update because received stale metadata",
+            response.getMetadataVersion(),
+            currentMetadataVersion);
+        return;
+      }
+      metadataVersion.set(response.getMetadataVersion());
+    }
+    ImmutableMap<String, String> globalDataServers =
+        response.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> endpoint.getValue().getEndpoint()));
+
+    ImmutableList<String> windmillServers =
+        response.getDirectPathEndpointsList().stream()
+            .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+            .collect(toImmutableList());
+
+    serverMappingUpdater.accept(
+        WindmillEndpoints.builder()
+            .setGlobalDataServers(globalDataServers)
+            .setWindmillServers(windmillServers)
+            .build());
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s]",
+        getMetadataVersion(), workerMetadataRequest.getHeader());
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> requestObserver() {
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  private synchronized long getMetadataVersion() {
+    return metadataVersion.get();
+  }
+
+  /**
+   * GetWorkerMetadata is a server streaming API. There is no request observer since the client side

Review Comment:
   It seems that a lot of the changes here are necessary because this is not a bi-directional stream.
   
   Beyond the code complexity of supporting this, I'm concerned about us not being able to send application-level heartbeat requests as we've had numerous issues in the past where sending was necessary to detect that the stream had died. See https://issues.apache.org/jira/browse/BEAM-10808 and https://github.com/apache/beam/pull/12688.
   
   I'm thinking that we should change the metadata streaming rpc to be bi-directional instead of just server side streaming to make sure we don't hit similar issues.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1315755710


##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,18 +758,15 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    // IPv6 address of a streaming engine windmill worker.
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;
   }
   repeated Endpoint work_endpoints = 2;
 
   // Maps from GlobalData tag to the endpoint that should be used for GetData
   // calls to retrieve that global data.
   map<string, Endpoint> global_data_endpoints = 3;
-
-  // DirectPath endpoints to be used by user workers for streaming engine jobs.
-  // DirectPath endpoints here are virtual IPv6 addresses of the windmill
-  // workers.
-  repeated Endpoint direct_path_endpoints = 4;

Review Comment:
   I don't see a newer commit.  You're likely not done addressing the other comments yet but just fyi in case you're waiting on me.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1316331506


##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,18 +758,15 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    // IPv6 address of a streaming engine windmill worker.
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;
   }
   repeated Endpoint work_endpoints = 2;
 
   // Maps from GlobalData tag to the endpoint that should be used for GetData
   // calls to retrieve that global data.
   map<string, Endpoint> global_data_endpoints = 3;
-
-  // DirectPath endpoints to be used by user workers for streaming engine jobs.
-  // DirectPath endpoints here are virtual IPv6 addresses of the windmill
-  // workers.
-  repeated Endpoint direct_path_endpoints = 4;

Review Comment:
   yep changes are still in the repo, was just using this to keep track of what i was finished with will tag you when finished thanks!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1709472015

   back to you @scwhittle thank you!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302326389


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getDirectPathEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue

Review Comment:
   confirmed it does https://chromium.googlesource.com/external/github.com/google/auto/+/auto-value-1.0/value/README.md#background



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1295273791


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+@AutoValue
+public abstract class WindmillEndpoints {
+  public abstract ImmutableMap<String, String> globalDataServers();

Review Comment:
   done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290777539


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+@AutoValue
+public abstract class WindmillEndpoints {
+  public abstract ImmutableMap<String, String> globalDataServers();
+
+  public abstract ImmutableList<String> windmillServers();

Review Comment:
   i guess we can convert the windmill server string to an object let me find the format



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1675418521

   waiting for changes to be finalized in windmill api proto before continuing to work on this @scwhittle 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302481067


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underlying stream observer that is being wrapped.
  */
-final class ForwardingClientResponseObserver<ReqT, RespT>
-    implements ClientResponseObserver<RespT, ReqT> {
+final class ForwardingClientResponseObserver<ResponseT, RequestT>

Review Comment:
   even after rebase this is showing up hmmm when i look at master https://github.com/apache/beam/blame/master/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java it still shows it as the old code?



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underlying stream observer that is being wrapped.
  */
-final class ForwardingClientResponseObserver<ReqT, RespT>
-    implements ClientResponseObserver<RespT, ReqT> {
+final class ForwardingClientResponseObserver<ResponseT, RequestT>

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308387141


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));

Review Comment:
   done



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;

Review Comment:
   done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308386236


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        endpointBuilder.setDirectEndpoint(endpointProto.getDirectEndpoint());
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+      return endpointBuilder.build();
+    }
+
+    public abstract Optional<String> directEndpoint();

Review Comment:
   added



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   i rebased this branch onto the branch that removes the sync stubs so i could reuse some of the code.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   changed method name, and added ClientInterceptor



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1317710815


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.locks.StampedLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+
+  /** Guards access to metadataVersion and latestResponse. */
+  private final StampedLock metadataLock;
+
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer,
+      StampedLock metadataLock) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+    this.metadataLock = metadataLock;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater,
+            new StampedLock());
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater,
+      StampedLock lock) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater,
+            lock);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse().
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long readStamp = metadataLock.tryOptimisticRead();

Review Comment:
   done.  using an object and synchronizing on it for read/write access



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1317747300


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,192 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+
+  private static final Logger LOG = LoggerFactory.getLogger(WindmillEndpoints.class);
+
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataEndpoints(globalDataServers)
+        .setWindmillEndpoints(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataEndpoints();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillEndpoints();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        parseDirectEndpoint(endpointProto.getDirectEndpoint())
+            .ifPresent(endpointBuilder::setDirectEndpoint);
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+
+      Endpoint endpoint = endpointBuilder.build();
+
+      if (!endpoint.directEndpoint().isPresent() && !endpoint.workerToken().isPresent()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "direct_endpoint=[%s] not present or could not be parsed, and worker_token"
+                    + " not present. At least one of these fields is required.",
+                endpointProto.getDirectEndpoint()));
+      }
+
+      return endpoint;
+    }
+
+    private static Optional<WindmillServiceAddress> parseDirectEndpoint(String directEndpoint) {
+      Optional<WindmillServiceAddress> directEndpointIpV6Address =
+          tryParseDirectEndpointIntoIpV6Address(directEndpoint).map(WindmillServiceAddress::create);
+
+      return directEndpointIpV6Address.isPresent()
+          ? directEndpointIpV6Address
+          : tryParseEndpointIntoHostAndPort(directEndpoint).map(WindmillServiceAddress::create);
+    }
+
+    private static Optional<HostAndPort> tryParseEndpointIntoHostAndPort(String directEndpoint) {
+      try {
+        return Optional.of(HostAndPort.fromString(directEndpoint));
+      } catch (IllegalArgumentException e) {
+        LOG.warn("{} cannot be parsed into a gcpServiceAddress", directEndpoint);
+        return Optional.empty();
+      }
+    }
+
+    private static Optional<Inet6Address> tryParseDirectEndpointIntoIpV6Address(
+        String directEndpoint) {
+      InetAddress directEndpointAddress = null;
+      try {
+        directEndpointAddress = Inet6Address.getByName(directEndpoint);
+      } catch (UnknownHostException e) {
+        LOG.warn(
+            "Error occurred trying to parse direct_endpoint={} into IPv6 address. Exception={}",
+            directEndpoint,
+            e.toString());
+      }
+
+      // Inet6Address.getByAddress returns either an IPv4 or an IPv6 address depending on the format
+      // of the direct_endpoint string.
+      if (!(directEndpointAddress instanceof Inet6Address)) {
+        LOG.warn(
+            "{} is not an IPv6 address. Direct endpoints are expected to be in IPv6 format.",
+            directEndpoint);
+        return Optional.empty();
+      }
+
+      return Optional.ofNullable((Inet6Address) directEndpointAddress);
+    }
+
+    public abstract Optional<WindmillServiceAddress> directEndpoint();

Review Comment:
   done. added comments and linked to proto and java files.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1289963961


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+@AutoValue
+public abstract class WindmillEndpoints {
+  public abstract ImmutableMap<String, String> globalDataServers();
+
+  public abstract ImmutableList<String> windmillServers();

Review Comment:
   add a comment about what string is
   
   alternatively should it be some sort of endpoint object instead of a string?



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -130,16 +127,18 @@ private static long debugDuration(long nowMs, long startMs) {
   protected abstract void onNewStream();
 
   /** Returns whether there are any pending requests that should be retried on a stream break. */
-  protected abstract boolean hasPendingRequests();
+  protected boolean hasPendingRequests() {

Review Comment:
   revert? Since this is related to correctness a default seems dangerous since it makes it possible to skip implementing correctly.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+@AutoValue
+public abstract class WindmillEndpoints {
+  public abstract ImmutableMap<String, String> globalDataServers();

Review Comment:
   Comment on what key and value are



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcCommitWorkStream.java:
##########
@@ -56,7 +56,8 @@ final class GrpcCommitWorkStream
   private final int streamingRpcBatchLimit;
 
   private GrpcCommitWorkStream(
-      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      Function<StreamObserver<StreamingCommitResponse>, StreamObserver<StreamingCommitWorkRequest>>
+          commitWorkRpc,

Review Comment:
   nit: commitWorkRpc is a confusing name to me for a function.  What about createCommitWorkRpcFn or startCommitWorkRpcFn etc?
   
   ditto for the other files



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -294,7 +293,10 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.size() == 1) {
       return stubList.get(0);
     }
-    return stubList.get(rand.nextInt(stubList.size()));
+    return stubList
+        .get(rand.nextInt(stubList.size()))
+        .withDeadlineAfter(
+            AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   this deadline is not being applied to stub returned above if only 1 stub



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,189 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        getWorkerMetadataRpc,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private long updateMetadataVersion(long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    metadataVersion.getAndUpdate(
+        current -> updateMetadataVersion(current, response.getMetadataVersion()));

Review Comment:
   if we get a version before our current version this seems to leave metadataVersion at the more current version *but* then it goes ahead and returns the stale response instead of filtering it.



##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,348 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String name = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataStream(

Review Comment:
   nit: add test to name? getWorkerMetadataTestStream



##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,348 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String name = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(name)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(InProcessChannelBuilder.forName(name).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(GLOBAL_DATA_ENDPOINTS.containsKey(key)));
+    assertEquals(
+        DIRECT_PATH_ENDPOINTS.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            DIRECT_PATH_ENDPOINTS.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(mockResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_consumesSubsequentResponseMetadata() {
+    WorkerMetadataResponse initialResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(initialResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+
+    List<WorkerMetadataResponse.Endpoint> newDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("newWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> newGlobalDataEndpoints = new HashMap<>();
+    newGlobalDataEndpoints.put(
+        "new_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("newGlobalData").build());
+
+    WorkerMetadataResponse newWorkMetadataResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(initialResponse.getMetadataVersion() + 1)
+            .addAllDirectPathEndpoints(newDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(newGlobalDataEndpoints)
+            .build();
+
+    testStub.injectWorkerMetadata(newWorkMetadataResponse);
+
+    assertEquals(
+        newGlobalDataEndpoints.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(newGlobalDataEndpoints.containsKey(key)));
+    assertEquals(
+        newDirectPathEndpoints.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            newDirectPathEndpoints.stream()
+                .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+                .collect(Collectors.toList())));
+    assertEquals(newWorkMetadataResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_doesNotConsumeResponseIfMetadataStale() {
+    WorkerMetadataResponse freshEndpoints =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(2)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(freshEndpoints);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    List<WorkerMetadataResponse.Endpoint> staleDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setEndpoint("staleWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> staleGlobalDataEndpoints = new HashMap<>();
+    staleGlobalDataEndpoints.put(
+        "stale_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setEndpoint("staleGlobalData").build());
+
+    testStub.injectWorkerMetadata(

Review Comment:
   see comment in code. It seems like this test should fail with the current logic.  It would also pass if this injection is somehow not reaching the stream so maybe some wiring is wrong?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290770865


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -130,16 +127,18 @@ private static long debugDuration(long nowMs, long startMs) {
   protected abstract void onNewStream();
 
   /** Returns whether there are any pending requests that should be retried on a stream break. */
-  protected abstract boolean hasPendingRequests();
+  protected boolean hasPendingRequests() {

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1707501972

   @scwhittle ready for another round of reviews!
   
   i had the WindmillServiceAddress.java file in the next CL in the chain, but I moved it to this one since it fits with the updates.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1295579721


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underlying stream observer that is being wrapped.
  */
-final class ForwardingClientResponseObserver<ReqT, RespT>
-    implements ClientResponseObserver<RespT, ReqT> {
+final class ForwardingClientResponseObserver<ResponseT, RequestT>

Review Comment:
   Can you rebase? I think a lot of this has already been merged with your previous PR



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getDirectPathEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue

Review Comment:
   just checking that Autovalue automatically creates equals/hash methods



##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,7 +758,8 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;

Review Comment:
   I think we want to remove direct_path_endpoints and just use work_endpoints



##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(GLOBAL_DATA_ENDPOINTS.containsKey(key)));
+    assertEquals(
+        DIRECT_PATH_ENDPOINTS.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            DIRECT_PATH_ENDPOINTS.stream()
+                .map(WindmillEndpoints.Endpoint::from)
+                .collect(Collectors.toList())));
+    assertEquals(mockResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_consumesSubsequentResponseMetadata() {
+    WorkerMetadataResponse initialResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(initialResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+
+    List<WorkerMetadataResponse.Endpoint> newDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setDirectEndpoint("newWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> newGlobalDataEndpoints = new HashMap<>();
+    newGlobalDataEndpoints.put(
+        "new_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("newGlobalData").build());
+
+    WorkerMetadataResponse newWorkMetadataResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(initialResponse.getMetadataVersion() + 1)
+            .addAllDirectPathEndpoints(newDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(newGlobalDataEndpoints)
+            .build();
+
+    testStub.injectWorkerMetadata(newWorkMetadataResponse);
+
+    assertEquals(
+        newGlobalDataEndpoints.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(newGlobalDataEndpoints.containsKey(key)));
+    assertEquals(
+        newDirectPathEndpoints.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            newDirectPathEndpoints.stream()
+                .map(WindmillEndpoints.Endpoint::from)
+                .collect(Collectors.toList())));
+    assertEquals(newWorkMetadataResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_doesNotConsumeResponseIfMetadataStale() {
+    WorkerMetadataResponse freshEndpoints =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(2)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(freshEndpoints);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    List<WorkerMetadataResponse.Endpoint> staleDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setDirectEndpoint("staleWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> staleGlobalDataEndpoints = new HashMap<>();
+    staleGlobalDataEndpoints.put(
+        "stale_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("staleGlobalData").build());
+
+    testStub.injectWorkerMetadata(
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(staleDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(staleGlobalDataEndpoints)
+            .build());
+
+    // Should have ignored the stale update and only used initial.
+    verify(testWindmillEndpointsConsumer).accept(WindmillEndpoints.from(freshEndpoints));
+    verifyNoMoreInteractions(testWindmillEndpointsConsumer);
+    assertEquals(freshEndpoints.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_correctlyAddsAndRemovesStreamFromRegistry() {
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(
+            WorkerMetadataResponse.newBuilder()
+                .setMetadataVersion(1)
+                .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+                .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+                .build());
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    stream =
+        getWorkerMetadataTestStream(
+            testStub, new AtomicLong(), new TestWindmillEndpointsConsumer());
+
+    assertTrue(streamRegistry.contains(stream));
+    stream.close();
+    assertFalse(streamRegistry.contains(stream));
+  }
+
+  @Test
+  public void testSendHealthCheck() {
+    TestGetWorkMetadataRequestObserver requestObserver =
+        Mockito.spy(
+            new TestGetWorkMetadataRequestObserver(
+                WorkerMetadataResponse.newBuilder()
+                    .setMetadataVersion(1)
+                    .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+                    .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+                    .build()));
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    stream =
+        getWorkerMetadataTestStream(
+            testStub, new AtomicLong(), new TestWindmillEndpointsConsumer());
+    stream.sendHealthCheck();
+
+    verify(requestObserver).onNext(WorkerMetadataRequest.getDefaultInstance());
+  }
+
+  private static class GetWorkerMetadataTestStub
+      extends CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase {
+    private final TestGetWorkMetadataRequestObserver requestObserver;
+    private @Nullable StreamObserver<WorkerMetadataResponse> responseObserver;
+
+    private GetWorkerMetadataTestStub(TestGetWorkMetadataRequestObserver requestObserver) {
+      this.requestObserver = requestObserver;
+    }
+
+    @Override
+    public StreamObserver<WorkerMetadataRequest> getWorkerMetadataStream(
+        StreamObserver<WorkerMetadataResponse> responseObserver) {
+      if (this.responseObserver == null) {
+        this.responseObserver = responseObserver;
+        requestObserver.responseObserver = this.responseObserver;
+      }
+
+      return requestObserver;
+    }
+
+    private void injectWorkerMetadata(WorkerMetadataResponse response) {
+      if (responseObserver != null) {
+        responseObserver.onNext(response);
+      }
+    }
+  }
+
+  private static class TestGetWorkMetadataRequestObserver
+      implements StreamObserver<WorkerMetadataRequest> {
+    private final WorkerMetadataResponse response;
+    private @Nullable StreamObserver<WorkerMetadataResponse> responseObserver;
+
+    private TestGetWorkMetadataRequestObserver(WorkerMetadataResponse response) {
+      this.response = response;
+    }
+
+    @Override
+    public void onNext(WorkerMetadataRequest workerMetadataRequest) {
+      responseObserver.onNext(response);

Review Comment:
   what about removing this way to send responses and just having tests explicitly call injectWorkerMetadata?
   I think it will be easier to see what tests are doing



##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);

Review Comment:
   nit: use some non-0 version, 0 might be correct if just default initialized on accident instead of propagated for example.
   
   (sorry if some of these are duplicated, I was putting comments on some other snapshot I think, not sure where they'll show up).



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,

Review Comment:
   integer? Can keep the AtomicLong internal for non-testing case to ensure it isn't modified externally



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1283886068


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,285 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final GetWorkMetadataStreamRequestObserver getWorkMetadataStreamRequestObserver;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        // This is a client unary request, server stream response RPC. There is no need to
+        // maintain the client side stream.
+        responseObserver -> null,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.getWorkMetadataStreamRequestObserver =
+        GetWorkMetadataStreamRequestObserver.newUnopenedStream(stub);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.fullyInitializeRequestObserver();
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.serverStream(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.fullyInitializeRequestObserver();
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Because {@link AbstractWindmillStream.ResponseObserver} is a non-static inner class, it (and
+   * its members) cannot be referenced until the constructor finishes executing and the object is
+   * completely instantiated. We gate this from being null by making the constructor private and
+   * exposing a static method so that an instance of {@link GrpcGetWorkerMetadataStream} will be
+   * created in the correct state.
+   */
+  private void fullyInitializeRequestObserver() {
+    getWorkMetadataStreamRequestObserver.init(
+        AbstractWindmillStream.ResponseObserver::new,
+        () -> {
+          synchronized (GrpcGetWorkerMetadataStream.this) {
+            streamClosed.set(true);
+          }
+        });
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> createNewRequestObserver() {
+    // we do not need to re-initiate the requestObserver here, since every
+    // GetWorkMetadataStreamRequestObserver.onNext() starts a new server stream.
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    synchronized (this) {
+      // If the metadataVersion newer or the same as the one in the response, the response data is
+      // stale, and we do not want to do anything.
+      long currentMetadataVersion = metadataVersion.get();
+      if (currentMetadataVersion >= response.getMetadataVersion()) {
+        LOG.info(
+            "Received metadata version={}; Current metadata version={}. "
+                + "Skipping update because received stale metadata",
+            response.getMetadataVersion(),
+            currentMetadataVersion);
+        return;
+      }
+      metadataVersion.set(response.getMetadataVersion());
+    }
+    ImmutableMap<String, String> globalDataServers =
+        response.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> endpoint.getValue().getEndpoint()));
+
+    ImmutableList<String> windmillServers =
+        response.getDirectPathEndpointsList().stream()
+            .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+            .collect(toImmutableList());
+
+    serverMappingUpdater.accept(
+        WindmillEndpoints.builder()
+            .setGlobalDataServers(globalDataServers)
+            .setWindmillServers(windmillServers)
+            .build());
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();

Review Comment:
   changed to bidi stream, so we will need this



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -73,9 +73,9 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final Executor executor;
   private final BackOff backoff;
   // Indicates if the current stream in requestObserver is closed by calling close() method
-  private final AtomicBoolean streamClosed;
+  protected AtomicBoolean streamClosed;
   private final AtomicLong startTimeMs;
-  private final AtomicLong lastSendTimeMs;
+  protected final AtomicLong lastSendTimeMs;

Review Comment:
   done.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -73,9 +73,9 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final Executor executor;
   private final BackOff backoff;
   // Indicates if the current stream in requestObserver is closed by calling close() method
-  private final AtomicBoolean streamClosed;
+  protected AtomicBoolean streamClosed;

Review Comment:
   done.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,285 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final GetWorkMetadataStreamRequestObserver getWorkMetadataStreamRequestObserver;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        // This is a client unary request, server stream response RPC. There is no need to
+        // maintain the client side stream.
+        responseObserver -> null,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.getWorkMetadataStreamRequestObserver =
+        GetWorkMetadataStreamRequestObserver.newUnopenedStream(stub);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.fullyInitializeRequestObserver();
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            stub,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.serverStream(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.fullyInitializeRequestObserver();
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Because {@link AbstractWindmillStream.ResponseObserver} is a non-static inner class, it (and
+   * its members) cannot be referenced until the constructor finishes executing and the object is
+   * completely instantiated. We gate this from being null by making the constructor private and
+   * exposing a static method so that an instance of {@link GrpcGetWorkerMetadataStream} will be
+   * created in the correct state.
+   */
+  private void fullyInitializeRequestObserver() {
+    getWorkMetadataStreamRequestObserver.init(
+        AbstractWindmillStream.ResponseObserver::new,
+        () -> {
+          synchronized (GrpcGetWorkerMetadataStream.this) {
+            streamClosed.set(true);
+          }
+        });
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> createNewRequestObserver() {
+    // we do not need to re-initiate the requestObserver here, since every
+    // GetWorkMetadataStreamRequestObserver.onNext() starts a new server stream.
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    synchronized (this) {
+      // If the metadataVersion newer or the same as the one in the response, the response data is
+      // stale, and we do not want to do anything.
+      long currentMetadataVersion = metadataVersion.get();
+      if (currentMetadataVersion >= response.getMetadataVersion()) {
+        LOG.info(
+            "Received metadata version={}; Current metadata version={}. "
+                + "Skipping update because received stale metadata",
+            response.getMetadataVersion(),
+            currentMetadataVersion);
+        return;
+      }
+      metadataVersion.set(response.getMetadataVersion());
+    }
+    ImmutableMap<String, String> globalDataServers =
+        response.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> endpoint.getValue().getEndpoint()));
+
+    ImmutableList<String> windmillServers =
+        response.getDirectPathEndpointsList().stream()
+            .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+            .collect(toImmutableList());
+
+    serverMappingUpdater.accept(
+        WindmillEndpoints.builder()
+            .setGlobalDataServers(globalDataServers)
+            .setWindmillServers(windmillServers)
+            .build());
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s]",
+        getMetadataVersion(), workerMetadataRequest.getHeader());
+  }
+
+  @Override
+  protected StreamObserver<WorkerMetadataRequest> requestObserver() {
+    return getWorkMetadataStreamRequestObserver;
+  }
+
+  private synchronized long getMetadataVersion() {
+    return metadataVersion.get();
+  }
+
+  /**
+   * GetWorkerMetadata is a server streaming API. There is no request observer since the client side

Review Comment:
   done. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308258594


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   changed method name, and added ClientInterceptor



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1307199364


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));

Review Comment:
   how about inlining the non-logging part of updateMetadataVersion here, and moving the logging below to else? Seems clearer than logging in the separate method.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));
+
+    if (updatedMetadataVersion > currentMetadataVersion) {
+      serverMappingUpdater.accept(WindmillEndpoints.from(response));
+    }
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected boolean hasPendingRequests() {
+    return false;
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void sendHealthCheck() {
+    send(HEALTH_CHECK_REQUEST);
+  }
+
+  @Override
+  protected void appendSpecificHtml(PrintWriter writer) {
+    writer.format(
+        "GetWorkerMetadataStream: version=[%d] , job_header=[%s]",

Review Comment:
   could show the latest response? could be useful.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            new AtomicLong(metadataVersion),
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private static long updateMetadataVersion(
+      long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    long currentMetadataVersion = metadataVersion.get();
+    long updatedMetadataVersion =
+        metadataVersion.updateAndGet(
+            current -> updateMetadataVersion(current, response.getMetadataVersion()));
+
+    if (updatedMetadataVersion > currentMetadataVersion) {
+      serverMappingUpdater.accept(WindmillEndpoints.from(response));

Review Comment:
   Do we need the atomic? I think that there is only one stream calling onResponse and thus those should be ordered. If we did have concurrent callers, this could be racy because you could have
   current version is 0
   T1: atomically updates to 1
   T2: atomically updates to 2, calls accept
   T1: resumes and calls accept
   then the consumer has last observed the payload for 1.
   
   So either:
   - remove the atomic if it isn't needed and document the thread-safety
   - or ensure that the ordering is maintained (ie could use synchronized and update version and call accept under it)



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataServers(globalDataServers)
+        .setWindmillServers(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataServers();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillServers();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        endpointBuilder.setDirectEndpoint(endpointProto.getDirectEndpoint());
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+      return endpointBuilder.build();
+    }
+
+    public abstract Optional<String> directEndpoint();

Review Comment:
   is there a better way to represent ipv6 address than String?



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,

Review Comment:
   still AtomicLong here



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   could name stubWithAbsoluteDeadlineSet to make it obvious to callers instead of a comment.



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;

Review Comment:
   nit: serverMappingConsumer or serverMappingObserver?
   Updater makes it sound a little like it is the other direction to me



##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcWindmillServer.java:
##########
@@ -295,10 +294,12 @@ private synchronized CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alph
     if (stubList.isEmpty()) {
       throw new RuntimeException("windmillServiceEndpoint has not been set");
     }
-    if (stubList.size() == 1) {
-      return stubList.get(0);
-    }
-    return stubList.get(rand.nextInt(stubList.size()));
+
+    CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub stub =
+        (stubList.size() == 1 ? stubList.get(0) : stubList.get(rand.nextInt(stubList.size())));
+
+    return stub.withDeadlineAfter(
+        AbstractWindmillStream.DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);

Review Comment:
   see other cl about withDeadlineAfter
   
   either use some other mechanism or document that result of stub() should just be used and not cached.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1308210008


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,

Review Comment:
   whoops misread the comment as for test only. done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302476191


##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);

Review Comment:
   done. no worries!



##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStreamTest.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class GrpcGetWorkerMetadataStreamTest {
+  private static final List<WorkerMetadataResponse.Endpoint> DIRECT_PATH_ENDPOINTS =
+      Lists.newArrayList(
+          WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("direct_path").build());
+  private static final Map<String, WorkerMetadataResponse.Endpoint> GLOBAL_DATA_ENDPOINTS =
+      Maps.newHashMap();
+  private static final JobHeader TEST_JOB_HEADER =
+      JobHeader.newBuilder()
+          .setJobId("test_job")
+          .setWorkerId("test_worker")
+          .setProjectId("test_project")
+          .build();
+  private static final String FAKE_SERVER_NAME = "Fake server for GrpcGetWorkerMetadataStreamTest";
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+  private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
+  private final Set<AbstractWindmillStream<?, ?>> streamRegistry = new HashSet<>();
+  private ManagedChannel inProcessChannel;
+  private GrpcGetWorkerMetadataStream stream;
+
+  private GrpcGetWorkerMetadataStream getWorkerMetadataTestStream(
+      GetWorkerMetadataTestStub getWorkerMetadataTestStub,
+      AtomicLong metadataVersion,
+      Consumer<WindmillEndpoints> endpointsConsumer) {
+    serviceRegistry.addService(getWorkerMetadataTestStub);
+    return GrpcGetWorkerMetadataStream.forTesting(
+        responseObserver ->
+            CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel)
+                .getWorkerMetadataStream(responseObserver),
+        metadataVersion,
+        TEST_JOB_HEADER,
+        streamRegistry,
+        endpointsConsumer);
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Server server =
+        InProcessServerBuilder.forName(FAKE_SERVER_NAME)
+            .fallbackHandlerRegistry(serviceRegistry)
+            .directExecutor()
+            .build()
+            .start();
+    inProcessChannel =
+        grpcCleanup.register(
+            InProcessChannelBuilder.forName(FAKE_SERVER_NAME).directExecutor().build());
+    grpcCleanup.register(server);
+    grpcCleanup.register(inProcessChannel);
+    GLOBAL_DATA_ENDPOINTS.put(
+        "global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("global_data").build());
+  }
+
+  @After
+  public void cleanUp() {
+    inProcessChannel.shutdownNow();
+  }
+
+  @Test
+  public void testGetWorkerMetadata() {
+    WorkerMetadataResponse mockResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        new TestWindmillEndpointsConsumer();
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(mockResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    assertEquals(
+        GLOBAL_DATA_ENDPOINTS.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(GLOBAL_DATA_ENDPOINTS.containsKey(key)));
+    assertEquals(
+        DIRECT_PATH_ENDPOINTS.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            DIRECT_PATH_ENDPOINTS.stream()
+                .map(WindmillEndpoints.Endpoint::from)
+                .collect(Collectors.toList())));
+    assertEquals(mockResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_consumesSubsequentResponseMetadata() {
+    WorkerMetadataResponse initialResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(initialResponse);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+
+    List<WorkerMetadataResponse.Endpoint> newDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setDirectEndpoint("newWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> newGlobalDataEndpoints = new HashMap<>();
+    newGlobalDataEndpoints.put(
+        "new_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("newGlobalData").build());
+
+    WorkerMetadataResponse newWorkMetadataResponse =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(initialResponse.getMetadataVersion() + 1)
+            .addAllDirectPathEndpoints(newDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(newGlobalDataEndpoints)
+            .build();
+
+    testStub.injectWorkerMetadata(newWorkMetadataResponse);
+
+    assertEquals(
+        newGlobalDataEndpoints.size(), testWindmillEndpointsConsumer.globalDataEndpoints.size());
+    testWindmillEndpointsConsumer.globalDataEndpoints.forEach(
+        (key, value) -> assertTrue(newGlobalDataEndpoints.containsKey(key)));
+    assertEquals(
+        newDirectPathEndpoints.size(), testWindmillEndpointsConsumer.windmillEndpoints.size());
+    assertTrue(
+        testWindmillEndpointsConsumer.windmillEndpoints.containsAll(
+            newDirectPathEndpoints.stream()
+                .map(WindmillEndpoints.Endpoint::from)
+                .collect(Collectors.toList())));
+    assertEquals(newWorkMetadataResponse.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_doesNotConsumeResponseIfMetadataStale() {
+    WorkerMetadataResponse freshEndpoints =
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(2)
+            .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+            .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+            .build();
+    TestWindmillEndpointsConsumer testWindmillEndpointsConsumer =
+        Mockito.spy(new TestWindmillEndpointsConsumer());
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(freshEndpoints);
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    AtomicLong metadataVersion = new AtomicLong(0);
+    stream = getWorkerMetadataTestStream(testStub, metadataVersion, testWindmillEndpointsConsumer);
+    List<WorkerMetadataResponse.Endpoint> staleDirectPathEndpoints =
+        Lists.newArrayList(
+            WorkerMetadataResponse.Endpoint.newBuilder()
+                .setDirectEndpoint("staleWindmillEndpoint")
+                .build());
+    Map<String, WorkerMetadataResponse.Endpoint> staleGlobalDataEndpoints = new HashMap<>();
+    staleGlobalDataEndpoints.put(
+        "stale_global_data",
+        WorkerMetadataResponse.Endpoint.newBuilder().setDirectEndpoint("staleGlobalData").build());
+
+    testStub.injectWorkerMetadata(
+        WorkerMetadataResponse.newBuilder()
+            .setMetadataVersion(1)
+            .addAllDirectPathEndpoints(staleDirectPathEndpoints)
+            .putAllGlobalDataEndpoints(staleGlobalDataEndpoints)
+            .build());
+
+    // Should have ignored the stale update and only used initial.
+    verify(testWindmillEndpointsConsumer).accept(WindmillEndpoints.from(freshEndpoints));
+    verifyNoMoreInteractions(testWindmillEndpointsConsumer);
+    assertEquals(freshEndpoints.getMetadataVersion(), metadataVersion.get());
+  }
+
+  @Test
+  public void testGetWorkerMetadata_correctlyAddsAndRemovesStreamFromRegistry() {
+    TestGetWorkMetadataRequestObserver requestObserver =
+        new TestGetWorkMetadataRequestObserver(
+            WorkerMetadataResponse.newBuilder()
+                .setMetadataVersion(1)
+                .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+                .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+                .build());
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    stream =
+        getWorkerMetadataTestStream(
+            testStub, new AtomicLong(), new TestWindmillEndpointsConsumer());
+
+    assertTrue(streamRegistry.contains(stream));
+    stream.close();
+    assertFalse(streamRegistry.contains(stream));
+  }
+
+  @Test
+  public void testSendHealthCheck() {
+    TestGetWorkMetadataRequestObserver requestObserver =
+        Mockito.spy(
+            new TestGetWorkMetadataRequestObserver(
+                WorkerMetadataResponse.newBuilder()
+                    .setMetadataVersion(1)
+                    .addAllDirectPathEndpoints(DIRECT_PATH_ENDPOINTS)
+                    .putAllGlobalDataEndpoints(GLOBAL_DATA_ENDPOINTS)
+                    .build()));
+    GetWorkerMetadataTestStub testStub = new GetWorkerMetadataTestStub(requestObserver);
+    stream =
+        getWorkerMetadataTestStream(
+            testStub, new AtomicLong(), new TestWindmillEndpointsConsumer());
+    stream.sendHealthCheck();
+
+    verify(requestObserver).onNext(WorkerMetadataRequest.getDefaultInstance());
+  }
+
+  private static class GetWorkerMetadataTestStub
+      extends CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase {
+    private final TestGetWorkMetadataRequestObserver requestObserver;
+    private @Nullable StreamObserver<WorkerMetadataResponse> responseObserver;
+
+    private GetWorkerMetadataTestStub(TestGetWorkMetadataRequestObserver requestObserver) {
+      this.requestObserver = requestObserver;
+    }
+
+    @Override
+    public StreamObserver<WorkerMetadataRequest> getWorkerMetadataStream(
+        StreamObserver<WorkerMetadataResponse> responseObserver) {
+      if (this.responseObserver == null) {
+        this.responseObserver = responseObserver;
+        requestObserver.responseObserver = this.responseObserver;
+      }
+
+      return requestObserver;
+    }
+
+    private void injectWorkerMetadata(WorkerMetadataResponse response) {
+      if (responseObserver != null) {
+        responseObserver.onNext(response);
+      }
+    }
+  }
+
+  private static class TestGetWorkMetadataRequestObserver
+      implements StreamObserver<WorkerMetadataRequest> {
+    private final WorkerMetadataResponse response;
+    private @Nullable StreamObserver<WorkerMetadataResponse> responseObserver;
+
+    private TestGetWorkMetadataRequestObserver(WorkerMetadataResponse response) {
+      this.response = response;
+    }
+
+    @Override
+    public void onNext(WorkerMetadataRequest workerMetadataRequest) {
+      responseObserver.onNext(response);

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1302477898


##########
runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto:
##########
@@ -758,7 +758,8 @@ message WorkerMetadataResponse {
   // CommitWorkStream. Each response on this stream replaces the previous, and
   // connections to endpoints that are no longer present should be closed.
   message Endpoint {
-    optional string endpoint = 1;
+    optional string direct_endpoint = 1;
+    optional string worker_token = 2;

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1283923064


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,190 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        getWorkerMetadataRpc,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.serverStream(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private void updateMetadataVersionIfStale(long responseMetadataVersion) {
+    long currentMetadataVersion;
+    synchronized (this) {
+      currentMetadataVersion = metadataVersion.get();
+      if (currentMetadataVersion < responseMetadataVersion) {
+        metadataVersion.set(responseMetadataVersion);
+        return;
+      }
+    }
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    updateMetadataVersionIfStale(response.getMetadataVersion());
+    ImmutableMap<String, String> globalDataServers =
+        response.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> endpoint.getValue().getEndpoint()));
+
+    ImmutableList<String> windmillServers =
+        response.getDirectPathEndpointsList().stream()
+            .map(WorkerMetadataResponse.Endpoint::getEndpoint)
+            .collect(toImmutableList());
+
+    serverMappingUpdater.accept(
+        WindmillEndpoints.builder()
+            .setGlobalDataServers(globalDataServers)
+            .setWindmillServers(windmillServers)
+            .build());
+  }
+
+  @Override
+  protected synchronized void onNewStream() {
+    send(workerMetadataRequest);
+  }
+
+  @Override
+  protected void startThrottleTimer() {
+    getWorkerMetadataThrottleTimer.start();
+  }
+
+  @Override
+  protected void sendHealthCheck() {}

Review Comment:
   will add behavior here once we decide



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1283885922


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java:
##########
@@ -27,23 +27,23 @@
  * <p>Used to wrap existing {@link StreamObserver}s to be able to install an {@link
  * ClientCallStreamObserver#setOnReadyHandler(Runnable) onReadyHandler}.
  *
- * <p>This is as thread-safe as the undering stream observer that is being wrapped.
+ * <p>This is as thread-safe as the underling stream observer that is being wrapped.

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1290771660


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,189 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import java.io.PrintWriter;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+
+  private final AtomicLong metadataVersion;
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingUpdater;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    super(
+        getWorkerMetadataRpc,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingUpdater = serverMappingUpdater;
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      AtomicLong metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      AtomicLong metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  private long updateMetadataVersion(long currentMetadataVersion, long responseMetadataVersion) {
+    if (currentMetadataVersion < responseMetadataVersion) {
+      return responseMetadataVersion;
+    }
+
+    // If the currentMetadataVersion is greater than or equal to one in the response, the response
+    // data is stale, and we do not want to do anything.
+    LOG.info(
+        "Received metadata version={}; Current metadata version={}. "
+            + "Skipping update because received stale metadata",
+        responseMetadataVersion,
+        currentMetadataVersion);
+    return currentMetadataVersion;
+  }
+
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    metadataVersion.getAndUpdate(
+        current -> updateMetadataVersion(current, response.getMetadataVersion()));

Review Comment:
   done. tried to use the Atomic* non blocking.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1714742275

   @scwhittle thanks for reviewing!
   @lostluck I had to add a dependency to the gradle file to use Google Truth assertions in a unit test and wanted to make sure i did that correctly.  Other than that, ready to merge thank you!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1317159057


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java:
##########
@@ -0,0 +1,192 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.auto.value.AutoValue;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Value class for holding endpoints used for communicating with Windmill service. */
+@AutoValue
+public abstract class WindmillEndpoints {
+
+  private static final Logger LOG = LoggerFactory.getLogger(WindmillEndpoints.class);
+
+  public static WindmillEndpoints from(
+      Windmill.WorkerMetadataResponse workerMetadataResponseProto) {
+    ImmutableMap<String, WindmillEndpoints.Endpoint> globalDataServers =
+        workerMetadataResponseProto.getGlobalDataEndpointsMap().entrySet().stream()
+            .collect(
+                toImmutableMap(
+                    Map.Entry::getKey, // global data key
+                    endpoint -> WindmillEndpoints.Endpoint.from(endpoint.getValue())));
+
+    ImmutableList<WindmillEndpoints.Endpoint> windmillServers =
+        workerMetadataResponseProto.getWorkEndpointsList().stream()
+            .map(WindmillEndpoints.Endpoint::from)
+            .collect(toImmutableList());
+
+    return WindmillEndpoints.builder()
+        .setGlobalDataEndpoints(globalDataServers)
+        .setWindmillEndpoints(windmillServers)
+        .build();
+  }
+
+  public static WindmillEndpoints.Builder builder() {
+    return new AutoValue_WindmillEndpoints.Builder();
+  }
+
+  /**
+   * Used by GetData GlobalDataRequest(s) to support Beam side inputs. Returns a map where the key
+   * is a global data tag and the value is the endpoint where the data associated with the global
+   * data tag resides.
+   *
+   * @see <a href="https://beam.apache.org/documentation/programming-guide/#side-inputs">Beam Side
+   *     Inputs</a>
+   */
+  public abstract ImmutableMap<String, Endpoint> globalDataEndpoints();
+
+  /**
+   * Used by GetWork/GetData/CommitWork calls to send, receive, and commit work directly to/from
+   * Windmill servers. Returns a list of endpoints used to communicate with the corresponding
+   * Windmill servers.
+   */
+  public abstract ImmutableList<Endpoint> windmillEndpoints();
+
+  @AutoValue
+  public abstract static class Endpoint {
+    public static Endpoint.Builder builder() {
+      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
+    }
+
+    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+      Endpoint.Builder endpointBuilder = Endpoint.builder();
+      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {
+        parseDirectEndpoint(endpointProto.getDirectEndpoint())
+            .ifPresent(endpointBuilder::setDirectEndpoint);
+      }
+      if (endpointProto.hasWorkerToken() && !endpointProto.getWorkerToken().isEmpty()) {
+        endpointBuilder.setWorkerToken(endpointProto.getWorkerToken());
+      }
+
+      Endpoint endpoint = endpointBuilder.build();
+
+      if (!endpoint.directEndpoint().isPresent() && !endpoint.workerToken().isPresent()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "direct_endpoint=[%s] not present or could not be parsed, and worker_token"
+                    + " not present. At least one of these fields is required.",
+                endpointProto.getDirectEndpoint()));
+      }
+
+      return endpoint;
+    }
+
+    private static Optional<WindmillServiceAddress> parseDirectEndpoint(String directEndpoint) {
+      Optional<WindmillServiceAddress> directEndpointIpV6Address =
+          tryParseDirectEndpointIntoIpV6Address(directEndpoint).map(WindmillServiceAddress::create);
+
+      return directEndpointIpV6Address.isPresent()
+          ? directEndpointIpV6Address
+          : tryParseEndpointIntoHostAndPort(directEndpoint).map(WindmillServiceAddress::create);
+    }
+
+    private static Optional<HostAndPort> tryParseEndpointIntoHostAndPort(String directEndpoint) {
+      try {
+        return Optional.of(HostAndPort.fromString(directEndpoint));
+      } catch (IllegalArgumentException e) {
+        LOG.warn("{} cannot be parsed into a gcpServiceAddress", directEndpoint);
+        return Optional.empty();
+      }
+    }
+
+    private static Optional<Inet6Address> tryParseDirectEndpointIntoIpV6Address(
+        String directEndpoint) {
+      InetAddress directEndpointAddress = null;
+      try {
+        directEndpointAddress = Inet6Address.getByName(directEndpoint);
+      } catch (UnknownHostException e) {
+        LOG.warn(
+            "Error occurred trying to parse direct_endpoint={} into IPv6 address. Exception={}",
+            directEndpoint,
+            e.toString());
+      }
+
+      // Inet6Address.getByAddress returns either an IPv4 or an IPv6 address depending on the format
+      // of the direct_endpoint string.
+      if (!(directEndpointAddress instanceof Inet6Address)) {
+        LOG.warn(
+            "{} is not an IPv6 address. Direct endpoints are expected to be in IPv6 format.",
+            directEndpoint);
+        return Optional.empty();
+      }
+
+      return Optional.ofNullable((Inet6Address) directEndpointAddress);
+    }
+
+    public abstract Optional<WindmillServiceAddress> directEndpoint();

Review Comment:
   nit: move public accessors above all the private methods?
   add some comments about how these shoudl be used or refer to the proto file and add a comment there?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] scwhittle commented on pull request #27767: add getworkermetadata streaming rpc

Posted by "scwhittle (via GitHub)" <gi...@apache.org>.
scwhittle commented on PR #27767:
URL: https://github.com/apache/beam/pull/27767#issuecomment-1702395782

   It looks like you added a merge commit, instead could you rebase your change?
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1313250456


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/grpcclient/GrpcGetWorkerMetadataStream.java:
##########
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.grpcclient;
+
+import java.io.PrintWriter;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.AbstractWindmillStream;
+import org.apache.beam.runners.dataflow.worker.windmill.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkerMetadataResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillStream.GetWorkerMetadataStream;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.vendor.grpc.v1p54p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+final class GrpcGetWorkerMetadataStream
+    extends AbstractWindmillStream<WorkerMetadataRequest, WorkerMetadataResponse>
+    implements GetWorkerMetadataStream {
+  private static final Logger LOG = LoggerFactory.getLogger(GrpcGetWorkerMetadataStream.class);
+  private static final WorkerMetadataRequest HEALTH_CHECK_REQUEST =
+      WorkerMetadataRequest.getDefaultInstance();
+  private final WorkerMetadataRequest workerMetadataRequest;
+  private final ThrottleTimer getWorkerMetadataThrottleTimer;
+  private final Consumer<WindmillEndpoints> serverMappingConsumer;
+  private long metadataVersion;
+  private WorkerMetadataResponse latestResponse;
+
+  private GrpcGetWorkerMetadataStream(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      long metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingConsumer) {
+    super(
+        startGetWorkerMetadataRpcFn,
+        backoff,
+        streamObserverFactory,
+        streamRegistry,
+        logEveryNStreamFailures);
+    this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
+    this.metadataVersion = metadataVersion;
+    this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
+    this.serverMappingConsumer = serverMappingConsumer;
+    this.latestResponse = WorkerMetadataResponse.getDefaultInstance();
+  }
+
+  public static GrpcGetWorkerMetadataStream create(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          startGetWorkerMetadataRpcFn,
+      BackOff backoff,
+      StreamObserverFactory streamObserverFactory,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      int logEveryNStreamFailures,
+      JobHeader jobHeader,
+      int metadataVersion,
+      ThrottleTimer getWorkerMetadataThrottleTimer,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream getWorkerMetadataStream =
+        new GrpcGetWorkerMetadataStream(
+            startGetWorkerMetadataRpcFn,
+            backoff,
+            streamObserverFactory,
+            streamRegistry,
+            logEveryNStreamFailures,
+            jobHeader,
+            metadataVersion,
+            getWorkerMetadataThrottleTimer,
+            serverMappingUpdater);
+    getWorkerMetadataStream.startStream();
+    return getWorkerMetadataStream;
+  }
+
+  @VisibleForTesting
+  static GrpcGetWorkerMetadataStream forTesting(
+      Function<StreamObserver<WorkerMetadataResponse>, StreamObserver<WorkerMetadataRequest>>
+          getWorkerMetadataRpc,
+      int metadataVersion,
+      JobHeader jobHeader,
+      Set<AbstractWindmillStream<?, ?>> streamRegistry,
+      Consumer<WindmillEndpoints> serverMappingUpdater) {
+    GrpcGetWorkerMetadataStream testStream =
+        new GrpcGetWorkerMetadataStream(
+            getWorkerMetadataRpc,
+            FluentBackoff.DEFAULT.backoff(),
+            StreamObserverFactory.direct(DEFAULT_STREAM_RPC_DEADLINE_SECONDS * 2, 1),
+            streamRegistry,
+            1, // logEveryNStreamFailures
+            jobHeader,
+            metadataVersion,
+            new ThrottleTimer(),
+            serverMappingUpdater);
+    testStream.startStream();
+    return testStream;
+  }
+
+  /**
+   * Each instance of {@link AbstractWindmillStream} owns its own responseObserver that calls
+   * onResponse(). This method does not and should not ever have any concurrent callers.
+   */
+  @Override
+  protected void onResponse(WorkerMetadataResponse response) {
+    if (response.getMetadataVersion() > metadataVersion) {
+      metadataVersion = response.getMetadataVersion();
+      serverMappingConsumer.accept(WindmillEndpoints.from(response));
+    } else {
+      // If the currentMetadataVersion is greater than or equal to one in the response, the response
+      // data is stale, and we do not want to do anything.
+      LOG.info(
+          "Received metadata version={}; Current metadata version={}. "
+              + "Skipping update because received stale metadata",
+          response.getMetadataVersion(),
+          metadataVersion);
+    }
+
+    this.latestResponse = response;

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [beam] m-trieu commented on a diff in pull request #27767: add getworkermetadata streaming rpc

Posted by "m-trieu (via GitHub)" <gi...@apache.org>.
m-trieu commented on code in PR #27767:
URL: https://github.com/apache/beam/pull/27767#discussion_r1291778054


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/AbstractWindmillStream.java:
##########
@@ -133,13 +130,13 @@ private static long debugDuration(long nowMs, long startMs) {
   protected abstract boolean hasPendingRequests();
 
   /**
-   * Called when the stream is throttled due to resource exhausted errors. Will be called for each
-   * resource exhausted error not just the first. onResponse() must stop throttling on receipt of
-   * the first good message.
+   * Called when the client side stream is throttled due to resource exhausted errors. Will be
+   * called for each resource exhausted error not just the first. onResponse() must stop throttling
+   * on receipt of the first good message.
    */
   protected abstract void startThrottleTimer();
 
-  private StreamObserver<RequestT> requestObserver() {
+  protected StreamObserver<RequestT> requestObserver() {

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org