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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
