add interfaces for direct path, and StreamingEngineClient

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillConnection.java

@@ -0,0 +1,57 @@
+/*
+ * 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 java.util.Optional;
+import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints.Endpoint;
+import org.apache.beam.sdk.annotations.Internal;
+
+@AutoValue
+@Internal
+public abstract class WindmillConnection {
+  public static WindmillConnection from(
+      Endpoint windmillEndpoint,
+      Function<Endpoint, CloudWindmillServiceV1Alpha1Stub> endpointToStubFn) {
+    WindmillConnection.Builder windmillWorkerConnection = WindmillConnection.builder();
+
+    windmillEndpoint.workerToken().ifPresent(windmillWorkerConnection::setBackendWorkerToken);
+    windmillWorkerConnection.setStub(endpointToStubFn.apply(windmillEndpoint));
+
+    return windmillWorkerConnection.build();
+  }
+
+  public static Builder builder() {
+    return new AutoValue_WindmillConnection.Builder();
+  }
+
+  public abstract Optional<String> backendWorkerToken();
+
+  public abstract CloudWindmillServiceV1Alpha1Stub stub();
+
+  @AutoValue.Builder
+  abstract static class Builder {
+    abstract Builder setBackendWorkerToken(String backendWorkerToken);
+
+    abstract Builder setStub(CloudWindmillServiceV1Alpha1Stub stub);
+
+    abstract WindmillConnection build();
+  }
+}

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkTimingInfosTracker.java

@@ -35,8 +35,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-class GetWorkTimingInfosTracker {
-
+final class GetWorkTimingInfosTracker {
   private static final Logger LOG = LoggerFactory.getLogger(GetWorkTimingInfosTracker.class);
 
   private final Map<State, SumAndMaxDurations> aggregatedGetWorkStreamLatencies;
@@ -53,7 +52,7 @@ class GetWorkTimingInfosTracker {
     workItemCreationLatency = null;
   }
 
-  public void addTimingInfo(Collection<GetWorkStreamTimingInfo> infos) {
+  void addTimingInfo(Collection<GetWorkStreamTimingInfo> infos) {
     // We want to record duration for each stage and also be reflective on total work item
     // processing time. It can be tricky because timings of different
     // StreamingGetWorkResponseChunks can be interleaved. Current strategy is to record the
@@ -170,19 +169,18 @@ List<LatencyAttribution> getLatencyAttributions() {
     return latencyAttributions;
   }
 
-  public void reset() {
+  void reset() {
     this.aggregatedGetWorkStreamLatencies.clear();
     this.workItemCreationEndTime = Instant.EPOCH;
     this.workItemLastChunkReceivedByWorkerTime = Instant.EPOCH;
     this.workItemCreationLatency = null;
   }
 
   private static class SumAndMaxDurations {
-
     private Duration sum;
     private Duration max;
 
-    public SumAndMaxDurations(Duration sum, Duration max) {
+    private SumAndMaxDurations(Duration sum, Duration max) {
       this.sum = sum;
       this.max = max;
     }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java

@@ -43,7 +43,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-final class GrpcCommitWorkStream
+public final class GrpcCommitWorkStream
     extends AbstractWindmillStream<StreamingCommitWorkRequest, StreamingCommitResponse>
     implements CommitWorkStream {
   private static final Logger LOG = LoggerFactory.getLogger(GrpcCommitWorkStream.class);
@@ -82,7 +82,7 @@ private GrpcCommitWorkStream(
     this.streamingRpcBatchLimit = streamingRpcBatchLimit;
   }
 
-  static GrpcCommitWorkStream create(
+  public static GrpcCommitWorkStream create(
       Function<StreamObserver<StreamingCommitResponse>, StreamObserver<StreamingCommitWorkRequest>>
           startCommitWorkRpcFn,
       BackOff backoff,