From bf8106c491534a13f6a3e264584ef785db50b26c Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Thu, 9 Nov 2023 10:21:29 -0800 Subject: [PATCH 1/8] WIP: Implement CacheSerializer and providers --- .../io/requestresponse/CacheSerializer.java | 35 ++++++ .../CacheSerializerProvider.java | 34 ++++++ .../CacheSerializerProviders.java | 44 +++++++ ...rministicCoderCacheSerializerProvider.java | 20 ++++ .../JsonCacheSerializerProvider.java | 107 ++++++++++++++++++ 5 files changed, 240 insertions(+) create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java new file mode 100644 index 0000000000000..934dbfa294e89 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java @@ -0,0 +1,35 @@ +/* + * 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.io.requestresponse; + +import java.io.Serializable; + +/** + * Converts between a user defined type {@link T} and a byte array for cache storage. See {@link + * CacheSerializerProvider} for details on how to register a new {@link CacheSerializer} with {@link + * CacheSerializerProviders}. + */ +public interface CacheSerializer extends Serializable { + long serialVersionUID = 3183949171861894060L; + + /** Convert a byte array to a user defined type {@link T}. */ + T deserialize(byte[] bytes) throws UserCodeExecutionException; + + /** Convert a user defined type {@link T} to a byte array. */ + byte[] serialize(T t) throws UserCodeExecutionException; +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java new file mode 100644 index 0000000000000..e32b2a2826885 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java @@ -0,0 +1,34 @@ +/* + * 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.io.requestresponse; + +import com.google.auto.service.AutoService; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.io.Providers; + +/** + * Provides a {@link CacheSerializer} using a {@link Providers.Identifyable#identifier}. To register + * a new {@link CacheSerializer} with {@link CacheSerializerProviders}, one simply needs to add the + * {@link AutoService} annotation to the {@link CacheSerializer}'s associated {@link + * CacheSerializerProvider} and assign the {@link CacheSerializerProvider} class to {@link + * AutoService#value} See {@link JsonCacheSerializerProvider} for an example. + */ +public interface CacheSerializerProvider extends Providers.Identifyable { + Schema getConfigurationSchema(); + CacheSerializer getSerializer(Class clazz); +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java new file mode 100644 index 0000000000000..df0aba2376c11 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java @@ -0,0 +1,44 @@ +/* + * 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.io.requestresponse; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.Map; +import org.apache.beam.sdk.schemas.io.Providers; + +/** + * Provides the {@link CacheSerializer} using its associated {@link CacheSerializerProvider}. See + * {@link CacheSerializerProvider} for further information on the requirements to create and + * register a new {@link CacheSerializer} with the {@link CacheSerializerProviders}. + */ +public final class CacheSerializerProviders { + private CacheSerializerProviders() {} + + private static final Map PROVIDERS = + Providers.loadProviders(CacheSerializerProvider.class); + + /** Queries the service loader */ + public static CacheSerializer getSerializer(String id, Class clazz) { + CacheSerializerProvider provider = + checkStateNotNull( + PROVIDERS.get(id), "No serializer provider exists with identifier `%s`.", id); + + return provider.getSerializer(clazz); + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java new file mode 100644 index 0000000000000..5302fa793b50d --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java @@ -0,0 +1,20 @@ +package org.apache.beam.io.requestresponse; + +import org.apache.beam.sdk.schemas.Schema; + +public class DeterministicCoderCacheSerializerProvider implements CacheSerializerProvider { + @Override + public String identifier() { + return null; + } + + @Override + public Schema getConfigurationSchema() { + return null; + } + + @Override + public CacheSerializer getSerializer(Class clazz) { + return null; + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java new file mode 100644 index 0000000000000..f3dde071fe926 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java @@ -0,0 +1,107 @@ +/* + * 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.io.requestresponse; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.auto.service.AutoService; +import java.io.IOException; + +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.ProcessFunction; +import org.checkerframework.checker.nullness.qual.NonNull; + +/** + * An implementation of {@link CacheSerializerProvider} that provides a {@link CacheSerializer} + * responsible for converting between a user defined type to a byte array representation of a JSON + * string. + */ +@AutoService(CacheSerializerProvider.class) +public class JsonCacheSerializerProvider implements CacheSerializerProvider { + private static final String IDENTIFIER = "json"; + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + @Override + public String identifier() { + return IDENTIFIER; + } + + @Override + public Schema getConfigurationSchema() { + return null; + } + + @Override + public CacheSerializer getSerializer(Class clazz) { + if (!OBJECT_MAPPER.canSerialize(clazz)) { + throw new IllegalArgumentException(clazz + " cannot serialize to JSON"); + } + return new JsonCacheSerializer<>(new ToJsonFn<>(), new FromJsonFn<>(clazz)); + } + + private static class JsonCacheSerializer implements CacheSerializer { + + private final ToJsonFn serializeFn; + private final FromJsonFn deserializeFn; + + private JsonCacheSerializer(ToJsonFn serializeFn, FromJsonFn deserializeFn) { + this.serializeFn = serializeFn; + this.deserializeFn = deserializeFn; + } + + @Override + public T deserialize(byte[] bytes) throws UserCodeExecutionException { + return deserializeFn.apply(bytes); + } + + @Override + public byte[] serialize(T t) throws UserCodeExecutionException { + return serializeFn.apply(t); + } + } + + private static class ToJsonFn implements ProcessFunction<@NonNull T, byte[]> { + + @Override + public byte[] apply(@NonNull T input) throws UserCodeExecutionException { + try { + return OBJECT_MAPPER.writeValueAsBytes(input); + } catch (JsonProcessingException e) { + throw new UserCodeExecutionException(e); + } + } + } + + private static class FromJsonFn implements ProcessFunction { + + private final Class clazz; + + private FromJsonFn(Class clazz) { + this.clazz = clazz; + } + + @Override + public @NonNull T apply(byte[] input) throws UserCodeExecutionException { + try { + return OBJECT_MAPPER.readValue(input, clazz); + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + } + } +} From 9b72c316f2c8e6c3c2f9efaae34fdc2c833f0a1d Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Fri, 10 Nov 2023 01:07:42 +0000 Subject: [PATCH 2/8] wip --- sdks/java/io/rrio/build.gradle | 1 + .../beam/io/requestresponse/CacheRead.java | 124 ++++++++++++------ .../requestresponse/CacheReadUsingRedis.java | 76 +++++++++++ .../io/requestresponse/CacheSerializer.java | 33 +++-- .../CacheSerializerProvider.java | 34 ----- .../CacheSerializerProviders.java | 44 ------- .../DeterministicCoderCacheSerializer.java | 63 +++++++++ ...rministicCoderCacheSerializerProvider.java | 20 --- .../JsonCacheSerializerProvider.java | 107 --------------- .../beam/io/requestresponse/RedisClient.java | 10 ++ ...DeterministicCoderCacheSerializerTest.java | 5 + .../io/requestresponse/RedisClientTestIT.java | 22 ++++ 12 files changed, 287 insertions(+), 252 deletions(-) create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java diff --git a/sdks/java/io/rrio/build.gradle b/sdks/java/io/rrio/build.gradle index bfd030ce61dc5..4ecdf4e91df7a 100644 --- a/sdks/java/io/rrio/build.gradle +++ b/sdks/java/io/rrio/build.gradle @@ -50,6 +50,7 @@ dependencies { testImplementation platform(library.java.google_cloud_platform_libraries_bom) testImplementation library.java.google_http_client testImplementation library.java.junit + testImplementation library.java.hamcrest testImplementation library.java.testcontainers_base testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java index 3765d25370a66..dcad74789b435 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java @@ -18,83 +18,129 @@ package org.apache.beam.io.requestresponse; import com.google.auto.value.AutoValue; +import java.net.URI; import java.util.Map; import org.apache.beam.io.requestresponse.CacheRead.Result; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PInput; import org.apache.beam.sdk.values.POutput; import org.apache.beam.sdk.values.PValue; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; /** * {@link CacheRead} reads associated {@link ResponseT} types from {@link RequestT} types, if any - * exist. + * exist, returning a {@link Result} where {@link Result#getResponses}'s {@link KV} will have null + * {@link ResponseT} for {@link RequestT}s that have no associated {@link ResponseT} in the cache. */ -class CacheRead - extends PTransform, Result> { - - private static final TupleTag FAILURE_TAG = new TupleTag() {}; +class CacheRead<@NonNull RequestT, @Nullable ResponseT> + extends PTransform< + @NonNull PCollection<@NonNull RequestT>, + @NonNull Result<@NonNull RequestT, @Nullable ResponseT>> { + + static <@NonNull RequestT, @Nullable ResponseT> + CacheRead<@NonNull RequestT, @Nullable ResponseT> usingRedis( + URI uri, + CacheSerializer<@NonNull RequestT> requestSerializer, + CacheSerializer<@NonNull ResponseT> responseSerializer, + Coder<@NonNull RequestT> requestCoder, + Coder<@Nullable ResponseT> responseCoder) { + RedisClient client = new RedisClient(uri); + CacheReadUsingRedis<@NonNull RequestT, @Nullable ResponseT> caller = + new CacheReadUsingRedis<>(client, requestSerializer, responseSerializer); + return new CacheRead<>( + Configuration.<@NonNull RequestT, @Nullable ResponseT>builder() + .setCaller(caller) + .setSetupTeardown(caller) + .setResponseCoder(KvCoder.of(requestCoder, responseCoder)) + .build()); + } - // TODO(damondouglas): remove suppress warnings after instance utilized. - @SuppressWarnings({"unused"}) - private final Configuration configuration; + private final Configuration<@NonNull RequestT, @Nullable ResponseT> configuration; - private CacheRead(Configuration configuration) { + private CacheRead(Configuration<@NonNull RequestT, @Nullable ResponseT> configuration) { this.configuration = configuration; } + @Override + public @NonNull Result<@NonNull RequestT, @Nullable ResponseT> expand( + PCollection<@NonNull RequestT> input) { + Call.Result> result = + input.apply( + Call.of(configuration.getCaller(), configuration.getResponseCoder()) + .withSetupTeardown(configuration.getSetupTeardown())); + return new Result<>( + input.getPipeline(), + new TupleTag>() {}, + result.getResponses(), + new TupleTag() {}, + result.getFailures()); + } + /** Configuration details for {@link CacheRead}. */ @AutoValue - abstract static class Configuration { + abstract static class Configuration<@NonNull RequestT, @Nullable ResponseT> { - static Builder builder() { + static <@NonNull RequestT, @Nullable ResponseT> + Builder<@NonNull RequestT, @Nullable ResponseT> builder() { return new AutoValue_CacheRead_Configuration.Builder<>(); } - abstract Builder toBuilder(); + abstract Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> getCaller(); + + abstract SetupTeardown getSetupTeardown(); + + abstract Coder> getResponseCoder(); + + abstract Builder<@NonNull RequestT, @Nullable ResponseT> toBuilder(); @AutoValue.Builder - abstract static class Builder { + abstract static class Builder<@NonNull RequestT, @Nullable ResponseT> { - abstract Configuration build(); - } - } + abstract Builder setCaller( + Caller> value); - @Override - public Result expand(PCollection input) { - return Result.of( - new TupleTag>() {}, PCollectionTuple.empty(input.getPipeline())); + abstract Builder setSetupTeardown(SetupTeardown value); + + abstract Builder setResponseCoder(Coder> value); + + abstract Configuration<@NonNull RequestT, @Nullable ResponseT> build(); + } } /** - * The {@link Result} of reading RequestT {@link PCollection} elements yielding ResponseT {@link - * PCollection} elements. + * The {@link Result} of reading @NonNull RequestT {@link PCollection} elements yielding @Nullable + * ResponseT {@link PCollection} elements. */ - static class Result implements POutput { - - static Result of( - TupleTag> responseTag, PCollectionTuple pct) { - return new Result<>(responseTag, pct); - } + static class Result<@NonNull RequestT, @Nullable ResponseT> implements POutput { private final Pipeline pipeline; - private final TupleTag> responseTag; - private final PCollection> responses; + private final TupleTag> responseTag; + private final PCollection> responses; + private final TupleTag failureTag; private final PCollection failures; - private Result(TupleTag> responseTag, PCollectionTuple pct) { - this.pipeline = pct.getPipeline(); + private Result( + Pipeline pipeline, + TupleTag> responseTag, + PCollection> responses, + TupleTag failureTag, + PCollection failures) { + this.pipeline = pipeline; this.responseTag = responseTag; - this.responses = pct.get(responseTag); - this.failures = pct.get(FAILURE_TAG); + this.responses = responses; + this.failureTag = failureTag; + this.failures = failures; } - PCollection> getResponses() { + PCollection> getResponses() { return responses; } @@ -103,15 +149,15 @@ PCollection getFailures() { } @Override - public Pipeline getPipeline() { + public @NonNull Pipeline getPipeline() { return this.pipeline; } @Override - public Map, PValue> expand() { + public @NonNull Map<@NonNull TupleTag, @NonNull PValue> expand() { return ImmutableMap.of( responseTag, responses, - FAILURE_TAG, failures); + failureTag, failures); } @Override diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java new file mode 100644 index 0000000000000..e9d675da1ba01 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java @@ -0,0 +1,76 @@ +/* + * 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.io.requestresponse; + +import java.util.Base64; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.values.KV; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Reads an associated {@link ResponseT} from a {@link RequestT} using a {@link RedisClient}. + * Implements {@link Caller} and {@link SetupTeardown} so that {@link CacheReadUsingRedis} can be + * used by {@link Call}. + */ +class CacheReadUsingRedis<@NonNull RequestT, @Nullable ResponseT> + implements Caller<@NonNull RequestT, KV<@NonNull RequestT, ResponseT>>, SetupTeardown { + + private final RedisClient client; + private final CacheSerializer<@NonNull RequestT> requestSerializer; + private final CacheSerializer<@NonNull ResponseT> responseSerializer; + + CacheReadUsingRedis( + RedisClient client, + CacheSerializer<@NonNull RequestT> requestSerializer, + CacheSerializer<@NonNull ResponseT> responseSerializer) { + this.client = client; + this.requestSerializer = requestSerializer; + this.responseSerializer = responseSerializer; + } + + /** Implements {@link SetupTeardown#setup} to be called during a {@link DoFn.Setup}. */ + @Override + public void setup() throws UserCodeExecutionException { + client.setup(); + } + + /** Implements {@link SetupTeardown#teardown} to be called during a {@link DoFn.Teardown}. */ + @Override + public void teardown() throws UserCodeExecutionException { + client.teardown(); + } + + /** + * Implements {@link Caller#call} reading an associated {@link RequestT} to a {@link ResponseT}. + * If no association exists, returns a {@link KV} with a null value. + */ + @Override + public KV<@NonNull RequestT, @Nullable ResponseT> call(@NonNull RequestT request) + throws UserCodeExecutionException { + byte[] requestBytes = requestSerializer.serialize(request); + byte[] encodedRequestBytes = Base64.getEncoder().encode(requestBytes); + byte @Nullable [] encodedResponseBytes = client.getBytes(encodedRequestBytes); + if (encodedResponseBytes == null) { + return KV.of(request, null); + } + byte[] responseBytes = Base64.getDecoder().decode(encodedResponseBytes); + ResponseT response = responseSerializer.deserialize(responseBytes); + return KV.of(request, response); + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java index 934dbfa294e89..48bd7eab1eba7 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java @@ -18,18 +18,35 @@ package org.apache.beam.io.requestresponse; import java.io.Serializable; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.checkerframework.checker.nullness.qual.NonNull; -/** - * Converts between a user defined type {@link T} and a byte array for cache storage. See {@link - * CacheSerializerProvider} for details on how to register a new {@link CacheSerializer} with {@link - * CacheSerializerProviders}. - */ -public interface CacheSerializer extends Serializable { - long serialVersionUID = 3183949171861894060L; +/** Converts between a user defined type {@link T} and a byte array for cache storage. */ +public interface CacheSerializer<@NonNull T> extends Serializable { /** Convert a byte array to a user defined type {@link T}. */ + @NonNull T deserialize(byte[] bytes) throws UserCodeExecutionException; /** Convert a user defined type {@link T} to a byte array. */ - byte[] serialize(T t) throws UserCodeExecutionException; + byte[] serialize(@NonNull T t) throws UserCodeExecutionException; + + /** + * Instantiates a {@link CacheSerializer} that converts between a JSON byte array and a user + * defined type {@link T}. + */ + static <@NonNull T> CacheSerializer usingJson(Class<@NonNull T> clazz) { + return new JsonCacheSerializer<>(clazz); + } + + /** + * Instantiates a {@link CacheSerializer} that converts between a byte array and a user defined + * type {@link T}. Throws a {@link NonDeterministicException} if {@link Coder#verifyDeterministic} + * does not pass. + */ + static <@NonNull T> CacheSerializer<@NonNull T> usingDeterministicCoder( + Coder<@NonNull T> deterministicCoder) throws NonDeterministicException { + return new DeterministicCoderCacheSerializer<>(deterministicCoder); + } } diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java deleted file mode 100644 index e32b2a2826885..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProvider.java +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.service.AutoService; -import org.apache.beam.sdk.schemas.Schema; -import org.apache.beam.sdk.schemas.io.Providers; - -/** - * Provides a {@link CacheSerializer} using a {@link Providers.Identifyable#identifier}. To register - * a new {@link CacheSerializer} with {@link CacheSerializerProviders}, one simply needs to add the - * {@link AutoService} annotation to the {@link CacheSerializer}'s associated {@link - * CacheSerializerProvider} and assign the {@link CacheSerializerProvider} class to {@link - * AutoService#value} See {@link JsonCacheSerializerProvider} for an example. - */ -public interface CacheSerializerProvider extends Providers.Identifyable { - Schema getConfigurationSchema(); - CacheSerializer getSerializer(Class clazz); -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java deleted file mode 100644 index df0aba2376c11..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializerProviders.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * 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.io.requestresponse; - -import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; - -import java.util.Map; -import org.apache.beam.sdk.schemas.io.Providers; - -/** - * Provides the {@link CacheSerializer} using its associated {@link CacheSerializerProvider}. See - * {@link CacheSerializerProvider} for further information on the requirements to create and - * register a new {@link CacheSerializer} with the {@link CacheSerializerProviders}. - */ -public final class CacheSerializerProviders { - private CacheSerializerProviders() {} - - private static final Map PROVIDERS = - Providers.loadProviders(CacheSerializerProvider.class); - - /** Queries the service loader */ - public static CacheSerializer getSerializer(String id, Class clazz) { - CacheSerializerProvider provider = - checkStateNotNull( - PROVIDERS.get(id), "No serializer provider exists with identifier `%s`.", id); - - return provider.getSerializer(clazz); - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java new file mode 100644 index 0000000000000..3c57dcc6c0d08 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java @@ -0,0 +1,63 @@ +/* + * 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.io.requestresponse; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; +import org.checkerframework.checker.nullness.qual.NonNull; + +public class DeterministicCoderCacheSerializer implements CacheSerializer { + + private final Coder deterministicCoder; + + /** + * Instantiates a {@link DeterministicCoderCacheSerializer} checking whether {@link + * Coder#verifyDeterministic} passes. + */ + DeterministicCoderCacheSerializer(Coder deterministicCoder) throws NonDeterministicException { + deterministicCoder.verifyDeterministic(); + this.deterministicCoder = deterministicCoder; + } + + /** Converts from a byte array to a user defined type {@link T}. */ + @Override + public @NonNull T deserialize(byte[] bytes) throws UserCodeExecutionException { + try { + return checkStateNotNull(deterministicCoder.decode(ByteSource.wrap(bytes).openStream())); + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + } + + /** Converts to a byte array from a user defined type {@link T}. */ + @Override + public byte[] serialize(@NonNull T t) throws UserCodeExecutionException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try { + deterministicCoder.encode(t, baos); + return baos.toByteArray(); + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java deleted file mode 100644 index 5302fa793b50d..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerProvider.java +++ /dev/null @@ -1,20 +0,0 @@ -package org.apache.beam.io.requestresponse; - -import org.apache.beam.sdk.schemas.Schema; - -public class DeterministicCoderCacheSerializerProvider implements CacheSerializerProvider { - @Override - public String identifier() { - return null; - } - - @Override - public Schema getConfigurationSchema() { - return null; - } - - @Override - public CacheSerializer getSerializer(Class clazz) { - return null; - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java deleted file mode 100644 index f3dde071fe926..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/JsonCacheSerializerProvider.java +++ /dev/null @@ -1,107 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.auto.service.AutoService; -import java.io.IOException; - -import org.apache.beam.sdk.schemas.Schema; -import org.apache.beam.sdk.transforms.ProcessFunction; -import org.checkerframework.checker.nullness.qual.NonNull; - -/** - * An implementation of {@link CacheSerializerProvider} that provides a {@link CacheSerializer} - * responsible for converting between a user defined type to a byte array representation of a JSON - * string. - */ -@AutoService(CacheSerializerProvider.class) -public class JsonCacheSerializerProvider implements CacheSerializerProvider { - private static final String IDENTIFIER = "json"; - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - - @Override - public String identifier() { - return IDENTIFIER; - } - - @Override - public Schema getConfigurationSchema() { - return null; - } - - @Override - public CacheSerializer getSerializer(Class clazz) { - if (!OBJECT_MAPPER.canSerialize(clazz)) { - throw new IllegalArgumentException(clazz + " cannot serialize to JSON"); - } - return new JsonCacheSerializer<>(new ToJsonFn<>(), new FromJsonFn<>(clazz)); - } - - private static class JsonCacheSerializer implements CacheSerializer { - - private final ToJsonFn serializeFn; - private final FromJsonFn deserializeFn; - - private JsonCacheSerializer(ToJsonFn serializeFn, FromJsonFn deserializeFn) { - this.serializeFn = serializeFn; - this.deserializeFn = deserializeFn; - } - - @Override - public T deserialize(byte[] bytes) throws UserCodeExecutionException { - return deserializeFn.apply(bytes); - } - - @Override - public byte[] serialize(T t) throws UserCodeExecutionException { - return serializeFn.apply(t); - } - } - - private static class ToJsonFn implements ProcessFunction<@NonNull T, byte[]> { - - @Override - public byte[] apply(@NonNull T input) throws UserCodeExecutionException { - try { - return OBJECT_MAPPER.writeValueAsBytes(input); - } catch (JsonProcessingException e) { - throw new UserCodeExecutionException(e); - } - } - } - - private static class FromJsonFn implements ProcessFunction { - - private final Class clazz; - - private FromJsonFn(Class clazz) { - this.clazz = clazz; - } - - @Override - public @NonNull T apply(byte[] input) throws UserCodeExecutionException { - try { - return OBJECT_MAPPER.readValue(input, clazz); - } catch (IOException e) { - throw new UserCodeExecutionException(e); - } - } - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/RedisClient.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/RedisClient.java index a87f5c191e4b0..a347a18524131 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/RedisClient.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/RedisClient.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import redis.clients.jedis.JedisPooled; import redis.clients.jedis.exceptions.JedisException; @@ -61,6 +62,15 @@ long decr(String key) throws UserCodeExecutionException { } } + /** Get a byte array associated with a byte array key. Returns null if key does not exist. */ + byte @Nullable [] getBytes(byte[] key) throws UserCodeExecutionException { + try { + return getSafeClient().get(key); + } catch (JedisException e) { + throw new UserCodeExecutionException(e); + } + } + /** * Get the long value stored by the key. Yields zero when key does not exist, keeping consistency * with Redis convention. Consider using {@link #exists} to query key existance. diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java new file mode 100644 index 0000000000000..a53b74d86d7fb --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java @@ -0,0 +1,5 @@ +package org.apache.beam.io.requestresponse; + +public class DeterministicCoderCacheSerializerTest { + +} \ No newline at end of file diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java index 1fbb320a5f23c..16b83b1637afa 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java @@ -20,6 +20,8 @@ import static org.apache.beam.sdk.io.common.SchemaAwareJavaBeans.allPrimitiveDataTypes; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -206,4 +208,24 @@ public void setThenDecrThenIncr_yieldsExpectedValue() throws UserCodeExecutionEx public void givenKeyNotExists_getLong_yieldsZero() throws UserCodeExecutionException { assertEquals(0L, externalClients.getActualClient().getLong(UUID.randomUUID().toString())); } + + @Test + public void givenKeyNotExists_getBytes_yieldsNull() throws UserCodeExecutionException { + assertNull( + externalClients + .getActualClient() + .getBytes(UUID.randomUUID().toString().getBytes(StandardCharsets.UTF_8))); + } + + @Test + public void givenKeyExists_getBytes_yieldsValue() throws UserCodeExecutionException { + byte[] keyBytes = UUID.randomUUID().toString().getBytes(StandardCharsets.UTF_8); + String expected = UUID.randomUUID().toString(); + byte[] expectedBytes = expected.getBytes(StandardCharsets.UTF_8); + externalClients.getValidatingClient().set(keyBytes, expectedBytes); + byte[] actualBytes = externalClients.getActualClient().getBytes(keyBytes); + assertNotNull(actualBytes); + String actual = new String(actualBytes, StandardCharsets.UTF_8); + assertEquals(expected, actual); + } } From 286263ce1e6da92310063ec04d253db0b5399de6 Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Thu, 9 Nov 2023 23:19:28 -0800 Subject: [PATCH 3/8] Condense Throttle into one class --- .../apache/beam/io/requestresponse/Cache.java | 192 +++++++++++ .../beam/io/requestresponse/CacheRead.java | 167 ---------- .../requestresponse/CacheReadUsingRedis.java | 76 ----- .../io/requestresponse/CacheSerializer.java | 52 --- .../beam/io/requestresponse/CacheWrite.java | 119 ------- .../DeterministicCoderCacheSerializer.java | 63 ---- .../apache/beam/io/requestresponse/Quota.java | 65 ++++ .../beam/io/requestresponse/Throttle.java | 306 ++++++++++++++++++ .../io/requestresponse/ThrottleDequeue.java | 101 ------ .../io/requestresponse/ThrottleEnqueue.java | 61 ---- .../requestresponse/ThrottleRefreshQuota.java | 55 ---- ...DeterministicCoderCacheSerializerTest.java | 5 - 12 files changed, 563 insertions(+), 699 deletions(-) create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheWrite.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java create mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleDequeue.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleEnqueue.java delete mode 100644 sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleRefreshQuota.java delete mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java new file mode 100644 index 0000000000000..9760b91f13579 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.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.io.requestresponse; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; + +/** Transforms for reading and writing request/response associations to a cache. */ +final class Cache { + + /** + * Read {@link RequestT} {@link ResponseT} associations from a cache. The {@link KV} value is null + * when no association exists. + */ + static < + @NonNull RequestT, + @Nullable ResponseT, + CallerSetupTeardownT extends + Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> & SetupTeardown> + PTransform< + PCollection<@NonNull RequestT>, + Call.Result>> + read( + CallerSetupTeardownT implementsCallerSetupTeardown, + Coder<@NonNull RequestT> requestTCoder, + Coder<@NonNull ResponseT> responseTCoder) { + return Call.ofCallerAndSetupTeardown( + implementsCallerSetupTeardown, KvCoder.of(requestTCoder, responseTCoder)); + } + + /** Write a {@link RequestT} {@link ResponseT} association to a cache. */ + static < + @NonNull RequestT, + @NonNull ResponseT, + CallerSetupTeardownT extends + Caller< + KV<@NonNull RequestT, @NonNull ResponseT>, + KV<@NonNull RequestT, @NonNull ResponseT>> + & SetupTeardown> + PTransform< + PCollection>, + Call.Result>> + write( + CallerSetupTeardownT implementsCallerSetupTeardown, + KvCoder<@NonNull RequestT, @NonNull ResponseT> kvCoder) { + return Call.ofCallerAndSetupTeardown(implementsCallerSetupTeardown, kvCoder); + } + + /** Provides cache read and write support using a {@link RedisClient}. */ + static class UsingRedis { + private final Coder<@NonNull RequestT> requestTCoder; + private final Coder<@NonNull ResponseT> responseTCoder; + private final RedisClient client; + + /** + * Instantiates a {@link UsingRedis}. Throws a {@link Coder.NonDeterministicException} if either + * {@link RequestT} or {@link ResponseT} {@link Coder} fails to {@link + * Coder#verifyDeterministic}. + */ + UsingRedis( + Coder<@NonNull RequestT> requestTCoder, + Coder<@NonNull ResponseT> responseTCoder, + RedisClient client) + throws Coder.NonDeterministicException { + this.client = client; + requestTCoder.verifyDeterministic(); + responseTCoder.verifyDeterministic(); + this.requestTCoder = requestTCoder; + this.responseTCoder = responseTCoder; + } + + /** Instantiates {@link Read}. */ + Read<@NonNull RequestT, @Nullable ResponseT> read() { + return new Read<>(this); + } + + /** + * Instantiates {@link Write}. The {@link Duration} determines how long the associated {@link + * RequestT} and {@link ResponseT} lasts in the cache. + */ + Write<@NonNull RequestT, @NonNull ResponseT> write(Duration expiry) { + return new Write<>(expiry, this); + } + + /** Reads associated {@link RequestT} {@link ResponseT} using a {@link RedisClient}. */ + static class Read<@NonNull RequestT, @Nullable ResponseT> + implements Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>>, + SetupTeardown { + + private final UsingRedis spec; + + private Read(UsingRedis spec) { + this.spec = spec; + } + + @Override + public KV<@NonNull RequestT, @Nullable ResponseT> call(@NonNull RequestT request) + throws UserCodeExecutionException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try { + spec.requestTCoder.encode(request, baos); + byte[] encodedRequest = baos.toByteArray(); + byte[] encodedResponse = spec.client.getBytes(encodedRequest); + if (encodedResponse == null) { + return KV.of(request, null); + } + ResponseT response = + checkStateNotNull( + spec.responseTCoder.decode(ByteSource.wrap(encodedResponse).openStream())); + return KV.of(request, response); + } catch (IllegalStateException | IOException e) { + throw new UserCodeExecutionException(e); + } + } + + @Override + public void setup() throws UserCodeExecutionException { + spec.client.setup(); + } + + @Override + public void teardown() throws UserCodeExecutionException { + spec.client.teardown(); + } + } + } + + static class Write<@NonNull RequestT, @NonNull ResponseT> + implements Caller< + KV<@NonNull RequestT, @NonNull ResponseT>, KV<@NonNull RequestT, @NonNull ResponseT>>, + SetupTeardown { + private final Duration expiry; + + private final UsingRedis spec; + + private Write(Duration expiry, UsingRedis spec) { + this.expiry = expiry; + this.spec = spec; + } + + @Override + public KV<@NonNull RequestT, @NonNull ResponseT> call( + KV<@NonNull RequestT, @NonNull ResponseT> request) throws UserCodeExecutionException { + ByteArrayOutputStream keyStream = new ByteArrayOutputStream(); + ByteArrayOutputStream valueStream = new ByteArrayOutputStream(); + try { + spec.requestTCoder.encode(request.getKey(), keyStream); + spec.responseTCoder.encode(request.getValue(), valueStream); + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + spec.client.setex(keyStream.toByteArray(), valueStream.toByteArray(), expiry); + return request; + } + + @Override + public void setup() throws UserCodeExecutionException { + spec.client.setup(); + } + + @Override + public void teardown() throws UserCodeExecutionException { + spec.client.teardown(); + } + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java deleted file mode 100644 index dcad74789b435..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheRead.java +++ /dev/null @@ -1,167 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.value.AutoValue; -import java.net.URI; -import java.util.Map; -import org.apache.beam.io.requestresponse.CacheRead.Result; -import org.apache.beam.sdk.Pipeline; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PInput; -import org.apache.beam.sdk.values.POutput; -import org.apache.beam.sdk.values.PValue; -import org.apache.beam.sdk.values.TupleTag; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -import org.checkerframework.checker.nullness.qual.NonNull; -import org.checkerframework.checker.nullness.qual.Nullable; - -/** - * {@link CacheRead} reads associated {@link ResponseT} types from {@link RequestT} types, if any - * exist, returning a {@link Result} where {@link Result#getResponses}'s {@link KV} will have null - * {@link ResponseT} for {@link RequestT}s that have no associated {@link ResponseT} in the cache. - */ -class CacheRead<@NonNull RequestT, @Nullable ResponseT> - extends PTransform< - @NonNull PCollection<@NonNull RequestT>, - @NonNull Result<@NonNull RequestT, @Nullable ResponseT>> { - - static <@NonNull RequestT, @Nullable ResponseT> - CacheRead<@NonNull RequestT, @Nullable ResponseT> usingRedis( - URI uri, - CacheSerializer<@NonNull RequestT> requestSerializer, - CacheSerializer<@NonNull ResponseT> responseSerializer, - Coder<@NonNull RequestT> requestCoder, - Coder<@Nullable ResponseT> responseCoder) { - RedisClient client = new RedisClient(uri); - CacheReadUsingRedis<@NonNull RequestT, @Nullable ResponseT> caller = - new CacheReadUsingRedis<>(client, requestSerializer, responseSerializer); - return new CacheRead<>( - Configuration.<@NonNull RequestT, @Nullable ResponseT>builder() - .setCaller(caller) - .setSetupTeardown(caller) - .setResponseCoder(KvCoder.of(requestCoder, responseCoder)) - .build()); - } - - private final Configuration<@NonNull RequestT, @Nullable ResponseT> configuration; - - private CacheRead(Configuration<@NonNull RequestT, @Nullable ResponseT> configuration) { - this.configuration = configuration; - } - - @Override - public @NonNull Result<@NonNull RequestT, @Nullable ResponseT> expand( - PCollection<@NonNull RequestT> input) { - Call.Result> result = - input.apply( - Call.of(configuration.getCaller(), configuration.getResponseCoder()) - .withSetupTeardown(configuration.getSetupTeardown())); - return new Result<>( - input.getPipeline(), - new TupleTag>() {}, - result.getResponses(), - new TupleTag() {}, - result.getFailures()); - } - - /** Configuration details for {@link CacheRead}. */ - @AutoValue - abstract static class Configuration<@NonNull RequestT, @Nullable ResponseT> { - - static <@NonNull RequestT, @Nullable ResponseT> - Builder<@NonNull RequestT, @Nullable ResponseT> builder() { - return new AutoValue_CacheRead_Configuration.Builder<>(); - } - - abstract Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> getCaller(); - - abstract SetupTeardown getSetupTeardown(); - - abstract Coder> getResponseCoder(); - - abstract Builder<@NonNull RequestT, @Nullable ResponseT> toBuilder(); - - @AutoValue.Builder - abstract static class Builder<@NonNull RequestT, @Nullable ResponseT> { - - abstract Builder setCaller( - Caller> value); - - abstract Builder setSetupTeardown(SetupTeardown value); - - abstract Builder setResponseCoder(Coder> value); - - abstract Configuration<@NonNull RequestT, @Nullable ResponseT> build(); - } - } - - /** - * The {@link Result} of reading @NonNull RequestT {@link PCollection} elements yielding @Nullable - * ResponseT {@link PCollection} elements. - */ - static class Result<@NonNull RequestT, @Nullable ResponseT> implements POutput { - - private final Pipeline pipeline; - private final TupleTag> responseTag; - private final PCollection> responses; - private final TupleTag failureTag; - private final PCollection failures; - - private Result( - Pipeline pipeline, - TupleTag> responseTag, - PCollection> responses, - TupleTag failureTag, - PCollection failures) { - this.pipeline = pipeline; - this.responseTag = responseTag; - this.responses = responses; - this.failureTag = failureTag; - this.failures = failures; - } - - PCollection> getResponses() { - return responses; - } - - PCollection getFailures() { - return failures; - } - - @Override - public @NonNull Pipeline getPipeline() { - return this.pipeline; - } - - @Override - public @NonNull Map<@NonNull TupleTag, @NonNull PValue> expand() { - return ImmutableMap.of( - responseTag, responses, - failureTag, failures); - } - - @Override - public void finishSpecifyingOutput( - String transformName, PInput input, PTransform transform) {} - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java deleted file mode 100644 index e9d675da1ba01..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheReadUsingRedis.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * 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.io.requestresponse; - -import java.util.Base64; -import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.values.KV; -import org.checkerframework.checker.nullness.qual.NonNull; -import org.checkerframework.checker.nullness.qual.Nullable; - -/** - * Reads an associated {@link ResponseT} from a {@link RequestT} using a {@link RedisClient}. - * Implements {@link Caller} and {@link SetupTeardown} so that {@link CacheReadUsingRedis} can be - * used by {@link Call}. - */ -class CacheReadUsingRedis<@NonNull RequestT, @Nullable ResponseT> - implements Caller<@NonNull RequestT, KV<@NonNull RequestT, ResponseT>>, SetupTeardown { - - private final RedisClient client; - private final CacheSerializer<@NonNull RequestT> requestSerializer; - private final CacheSerializer<@NonNull ResponseT> responseSerializer; - - CacheReadUsingRedis( - RedisClient client, - CacheSerializer<@NonNull RequestT> requestSerializer, - CacheSerializer<@NonNull ResponseT> responseSerializer) { - this.client = client; - this.requestSerializer = requestSerializer; - this.responseSerializer = responseSerializer; - } - - /** Implements {@link SetupTeardown#setup} to be called during a {@link DoFn.Setup}. */ - @Override - public void setup() throws UserCodeExecutionException { - client.setup(); - } - - /** Implements {@link SetupTeardown#teardown} to be called during a {@link DoFn.Teardown}. */ - @Override - public void teardown() throws UserCodeExecutionException { - client.teardown(); - } - - /** - * Implements {@link Caller#call} reading an associated {@link RequestT} to a {@link ResponseT}. - * If no association exists, returns a {@link KV} with a null value. - */ - @Override - public KV<@NonNull RequestT, @Nullable ResponseT> call(@NonNull RequestT request) - throws UserCodeExecutionException { - byte[] requestBytes = requestSerializer.serialize(request); - byte[] encodedRequestBytes = Base64.getEncoder().encode(requestBytes); - byte @Nullable [] encodedResponseBytes = client.getBytes(encodedRequestBytes); - if (encodedResponseBytes == null) { - return KV.of(request, null); - } - byte[] responseBytes = Base64.getDecoder().decode(encodedResponseBytes); - ResponseT response = responseSerializer.deserialize(responseBytes); - return KV.of(request, response); - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java deleted file mode 100644 index 48bd7eab1eba7..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheSerializer.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * 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.io.requestresponse; - -import java.io.Serializable; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.Coder.NonDeterministicException; -import org.checkerframework.checker.nullness.qual.NonNull; - -/** Converts between a user defined type {@link T} and a byte array for cache storage. */ -public interface CacheSerializer<@NonNull T> extends Serializable { - - /** Convert a byte array to a user defined type {@link T}. */ - @NonNull - T deserialize(byte[] bytes) throws UserCodeExecutionException; - - /** Convert a user defined type {@link T} to a byte array. */ - byte[] serialize(@NonNull T t) throws UserCodeExecutionException; - - /** - * Instantiates a {@link CacheSerializer} that converts between a JSON byte array and a user - * defined type {@link T}. - */ - static <@NonNull T> CacheSerializer usingJson(Class<@NonNull T> clazz) { - return new JsonCacheSerializer<>(clazz); - } - - /** - * Instantiates a {@link CacheSerializer} that converts between a byte array and a user defined - * type {@link T}. Throws a {@link NonDeterministicException} if {@link Coder#verifyDeterministic} - * does not pass. - */ - static <@NonNull T> CacheSerializer<@NonNull T> usingDeterministicCoder( - Coder<@NonNull T> deterministicCoder) throws NonDeterministicException { - return new DeterministicCoderCacheSerializer<>(deterministicCoder); - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheWrite.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheWrite.java deleted file mode 100644 index 25249c3e41b42..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/CacheWrite.java +++ /dev/null @@ -1,119 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.value.AutoValue; -import java.util.Map; -import org.apache.beam.io.requestresponse.CacheWrite.Result; -import org.apache.beam.sdk.Pipeline; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionTuple; -import org.apache.beam.sdk.values.PInput; -import org.apache.beam.sdk.values.POutput; -import org.apache.beam.sdk.values.PValue; -import org.apache.beam.sdk.values.TupleTag; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; - -/** - * {@link CacheWrite} writes associated {@link RequestT} and {@link ResponseT} pairs to a cache. - * Using {@link RequestT} and {@link ResponseT}'s {@link org.apache.beam.sdk.coders.Coder}, this - * transform writes encoded representations of this association. - */ -class CacheWrite - extends PTransform>, Result> { - - private static final TupleTag FAILURE_TAG = new TupleTag() {}; - - // TODO(damondouglas): remove suppress warnings after configuration is used. - @SuppressWarnings({"unused"}) - private final Configuration configuration; - - private CacheWrite(Configuration configuration) { - this.configuration = configuration; - } - - /** Configuration details for {@link CacheWrite}. */ - @AutoValue - abstract static class Configuration { - - static Builder builder() { - return new AutoValue_CacheWrite_Configuration.Builder<>(); - } - - abstract Builder toBuilder(); - - @AutoValue.Builder - abstract static class Builder { - - abstract Configuration build(); - } - } - - @Override - public Result expand(PCollection> input) { - return Result.of( - new TupleTag>() {}, PCollectionTuple.empty(input.getPipeline())); - } - - /** The {@link Result} of writing a request/response {@link KV} {@link PCollection}. */ - static class Result implements POutput { - - static Result of( - TupleTag> responseTag, PCollectionTuple pct) { - return new Result<>(responseTag, pct); - } - - private final Pipeline pipeline; - private final TupleTag> responseTag; - private final PCollection> responses; - private final PCollection failures; - - private Result(TupleTag> responseTag, PCollectionTuple pct) { - this.pipeline = pct.getPipeline(); - this.responseTag = responseTag; - this.responses = pct.get(responseTag); - this.failures = pct.get(FAILURE_TAG); - } - - public PCollection> getResponses() { - return responses; - } - - public PCollection getFailures() { - return failures; - } - - @Override - public Pipeline getPipeline() { - return this.pipeline; - } - - @Override - public Map, PValue> expand() { - return ImmutableMap.of( - responseTag, responses, - FAILURE_TAG, failures); - } - - @Override - public void finishSpecifyingOutput( - String transformName, PInput input, PTransform transform) {} - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java deleted file mode 100644 index 3c57dcc6c0d08..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializer.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.io.requestresponse; - -import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; - -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.Coder.NonDeterministicException; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; -import org.checkerframework.checker.nullness.qual.NonNull; - -public class DeterministicCoderCacheSerializer implements CacheSerializer { - - private final Coder deterministicCoder; - - /** - * Instantiates a {@link DeterministicCoderCacheSerializer} checking whether {@link - * Coder#verifyDeterministic} passes. - */ - DeterministicCoderCacheSerializer(Coder deterministicCoder) throws NonDeterministicException { - deterministicCoder.verifyDeterministic(); - this.deterministicCoder = deterministicCoder; - } - - /** Converts from a byte array to a user defined type {@link T}. */ - @Override - public @NonNull T deserialize(byte[] bytes) throws UserCodeExecutionException { - try { - return checkStateNotNull(deterministicCoder.decode(ByteSource.wrap(bytes).openStream())); - } catch (IOException e) { - throw new UserCodeExecutionException(e); - } - } - - /** Converts to a byte array from a user defined type {@link T}. */ - @Override - public byte[] serialize(@NonNull T t) throws UserCodeExecutionException { - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - try { - deterministicCoder.encode(t, baos); - return baos.toByteArray(); - } catch (IOException e) { - throw new UserCodeExecutionException(e); - } - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java new file mode 100644 index 0000000000000..c4d2af113cebb --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java @@ -0,0 +1,65 @@ +/* + * 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.io.requestresponse; + +import java.io.Serializable; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.joda.time.Duration; + +/** + * A data class that expresses a quota. Web API providers typically define a quota as the number of + * requests per time interval. + */ +public class Quota implements Serializable { + private final long numRequests; + private final @NonNull Duration interval; + + public Quota(long numRequests, @NonNull Duration interval) { + this.numRequests = numRequests; + this.interval = interval; + } + + /** The number of allowed requests. */ + public long getNumRequests() { + return numRequests; + } + + /** The duration context within which to allow requests. */ + public @NonNull Duration getInterval() { + return interval; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + Quota quota = (Quota) o; + return numRequests == quota.numRequests && Objects.equal(interval, quota.interval); + } + + @Override + public int hashCode() { + return Objects.hashCode(numRequests, interval); + } + + @Override + public String toString() { + return "Quota{" + "numRequests=" + numRequests + ", interval=" + interval + '}'; + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java new file mode 100644 index 0000000000000..eac9aacb4a3d9 --- /dev/null +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java @@ -0,0 +1,306 @@ +/* + * 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.io.requestresponse; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.URI; +import java.util.Optional; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.PeriodicImpulse; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.joda.time.Duration; +import org.joda.time.Instant; + +class ThrottleWithExternalResource< + @NonNull T, + ReporterT extends Caller<@NonNull String, @NonNull Long> & SetupTeardown, + EnqueuerT extends Caller<@NonNull T, Void> & SetupTeardown, + DequeuerT extends Caller<@NonNull Instant, @NonNull T> & SetupTeardown, + RefresherT extends Caller<@NonNull Instant, Void> & SetupTeardown> + extends PTransform, Call.Result<@NonNull T>> { + + /** Instantiate a {@link ThrottleWithExternalResource} using a {@link RedisClient}. */ + static <@NonNull T> + ThrottleWithExternalResource< + @NonNull T, + RedisReporter, + RedisEnqueuer<@NonNull T>, + RedisDequeur<@NonNull T>, + RedisRefresher> + usingRedis( + URI uri, + String quotaIdentifier, + String queueKey, + Quota quota, + Coder<@NonNull T> coder) + throws Coder.NonDeterministicException { + return new ThrottleWithExternalResource< + @NonNull T, + RedisReporter, + RedisEnqueuer<@NonNull T>, + RedisDequeur<@NonNull T>, + RedisRefresher>( + quota, + quotaIdentifier, + coder, + new RedisReporter(uri), + new RedisEnqueuer<>(uri, queueKey, coder), + new RedisDequeur<>(uri, coder, queueKey), + new RedisRefresher(uri, quota, quotaIdentifier)); + } + + private final TupleTag<@NonNull T> outputTag = new TupleTag<@NonNull T>() {}; + private static final TupleTag FAILURE_TAG = new TupleTag() {}; + private static final Duration THROTTLE_INTERVAL = Duration.standardSeconds(1L); + + private final Quota quota; + private final String quotaIdentifier; + private final Coder coder; + private final ReporterT reporterT; + private final EnqueuerT enqueuerT; + private final DequeuerT dequeuerT; + private final RefresherT refresherT; + + ThrottleWithExternalResource( + Quota quota, + String quotaIdentifier, + Coder coder, + ReporterT reporterT, + EnqueuerT enqueuerT, + DequeuerT dequeuerT, + RefresherT refresherT) + throws Coder.NonDeterministicException { + this.quotaIdentifier = quotaIdentifier; + this.reporterT = reporterT; + coder.verifyDeterministic(); + this.quota = quota; + this.coder = coder; + this.enqueuerT = enqueuerT; + this.dequeuerT = dequeuerT; + this.refresherT = refresherT; + } + + @Override + public Call.Result<@NonNull T> expand(PCollection<@NonNull T> input) { + Pipeline pipeline = input.getPipeline(); + + // Refresh known quota to control the throttle rate. + Call.Result refreshResult = + pipeline + .apply("quota/impulse", PeriodicImpulse.create().withInterval(quota.getInterval())) + .apply("quota/refresh", getRefresher()); + + // Enqueue T elements. + Call.Result enqueuResult = input.apply("enqueue", getEnqueuer()); + + // Perform Throttle. + PCollectionTuple pct = + pipeline + .apply("throttle/impulse", PeriodicImpulse.create().withInterval(THROTTLE_INTERVAL)) + .apply( + "throttle/fn", + ParDo.of(new ThrottleFn(quotaIdentifier, dequeuerT, reporterT, outputTag)) + .withOutputTags(outputTag, TupleTagList.of(FAILURE_TAG))); + + PCollection errors = + PCollectionList.of(refreshResult.getFailures()) + .and(enqueuResult.getFailures()) + .and(pct.get(FAILURE_TAG)) + .apply("errors/flatten", Flatten.pCollections()); + + return Call.Result.of( + coder, + outputTag, + PCollectionTuple.of(outputTag, pct.get(outputTag)).and(FAILURE_TAG, errors)); + } + + private Call<@NonNull Instant, Void> getRefresher() { + return Call.ofCallerAndSetupTeardown(refresherT, VoidCoder.of()); + } + + private Call<@NonNull T, Void> getEnqueuer() { + return Call.ofCallerAndSetupTeardown(enqueuerT, VoidCoder.of()); + } + + private class ThrottleFn extends DoFn<@NonNull Instant, @NonNull T> { + private final String quotaIdentifier; + private final DequeuerT dequeuerT; + private final ReporterT reporterT; + private final TupleTag<@NonNull T> outputTag; + + private ThrottleFn( + String quotaIdentifier, + DequeuerT dequeuerT, + ReporterT reporterT, + TupleTag<@NonNull T> outputTag) { + this.quotaIdentifier = quotaIdentifier; + this.dequeuerT = dequeuerT; + this.reporterT = reporterT; + this.outputTag = outputTag; + } + + @ProcessElement + public void process(@Element @NonNull Instant instant, MultiOutputReceiver receiver) { + // Check for available quota. + try { + if (reporterT.call(quotaIdentifier) <= 0L) { + return; + } + + // Dequeue an element if quota available. + T element = dequeuerT.call(instant); + // Emit element. + receiver.get(outputTag).output(element); + } catch (UserCodeExecutionException e) { + receiver + .get(FAILURE_TAG) + .output( + ApiIOError.builder() + .setRequestAsJsonString(String.format("{\"request\": \"%s\"}", instant)) + .setMessage(Optional.ofNullable(e.getMessage()).orElse("")) + .setObservedTimestamp(Instant.now()) + .setStackTrace(Throwables.getStackTraceAsString(e)) + .build()); + } + } + + @Setup + public void setup() throws UserCodeExecutionException { + enqueuerT.setup(); + dequeuerT.setup(); + reporterT.setup(); + } + + @Teardown + public void teardown() throws UserCodeExecutionException { + enqueuerT.teardown(); + dequeuerT.teardown(); + reporterT.teardown(); + } + } + + private static class RedisReporter extends RedisSetupTeardown + implements Caller<@NonNull String, @NonNull Long> { + private RedisReporter(URI uri) { + super(new RedisClient(uri)); + } + + @Override + public @NonNull Long call(@NonNull String request) throws UserCodeExecutionException { + return client.getLong(request); + } + } + + private static class RedisEnqueuer<@NonNull T> extends RedisSetupTeardown + implements Caller<@NonNull T, Void> { + private final String key; + private final Coder coder; + + private RedisEnqueuer(URI uri, String key, Coder coder) { + super(new RedisClient(uri)); + this.key = key; + this.coder = coder; + } + + @Override + public Void call(@NonNull T request) throws UserCodeExecutionException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try { + coder.encode(request, baos); + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + client.rpush(key, baos.toByteArray()); + return null; + } + } + + private static class RedisDequeur<@NonNull T> extends RedisSetupTeardown + implements Caller<@NonNull Instant, @NonNull T> { + + private final Coder coder; + private final String key; + + private RedisDequeur(URI uri, Coder coder, String key) { + super(new RedisClient(uri)); + this.coder = coder; + this.key = key; + } + + @Override + public @NonNull T call(@NonNull Instant request) throws UserCodeExecutionException { + byte[] bytes = client.lpop(key); + try { + return coder.decode(ByteSource.wrap(bytes).openStream()); + + } catch (IOException e) { + throw new UserCodeExecutionException(e); + } + } + } + + private static class RedisRefresher extends RedisSetupTeardown + implements Caller<@NonNull Instant, Void> { + private final Quota quota; + private final String key; + + private RedisRefresher(URI uri, Quota quota, String key) { + super(new RedisClient(uri)); + this.quota = quota; + this.key = key; + } + + @Override + public Void call(@NonNull Instant request) throws UserCodeExecutionException { + client.setex(key, quota.getNumRequests(), quota.getInterval()); + return null; + } + } + + private abstract static class RedisSetupTeardown implements SetupTeardown { + protected final RedisClient client; + + private RedisSetupTeardown(RedisClient client) { + this.client = client; + } + + @Override + public void setup() throws UserCodeExecutionException { + client.setup(); + } + + @Override + public void teardown() throws UserCodeExecutionException { + client.teardown(); + } + } +} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleDequeue.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleDequeue.java deleted file mode 100644 index 085b13b5e1120..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleDequeue.java +++ /dev/null @@ -1,101 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.value.AutoValue; -import java.util.Map; -import org.apache.beam.io.requestresponse.ThrottleDequeue.Result; -import org.apache.beam.sdk.Pipeline; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionTuple; -import org.apache.beam.sdk.values.PInput; -import org.apache.beam.sdk.values.POutput; -import org.apache.beam.sdk.values.PValue; -import org.apache.beam.sdk.values.TupleTag; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -import org.joda.time.Instant; - -/** - * {@link ThrottleDequeue} dequeues {@link RequestT} elements at a fixed rate yielding a {@link - * Result} containing the dequeued {@link RequestT} {@link PCollection} and a {@link ApiIOError} - * {@link PCollection} of any errors. - */ -class ThrottleDequeue extends PTransform, Result> { - - private static final TupleTag FAILURE_TAG = new TupleTag() {}; - - // TODO(damondouglas): remove suppress warnings after instance utilized. - @SuppressWarnings({"unused"}) - private final Configuration configuration; - - private ThrottleDequeue(Configuration configuration) { - this.configuration = configuration; - } - - @Override - public Result expand(PCollection input) { - // TODO(damondouglas): expand in a future PR. - return new Result<>(new TupleTag() {}, PCollectionTuple.empty(input.getPipeline())); - } - - @AutoValue - abstract static class Configuration { - - @AutoValue.Builder - abstract static class Builder { - abstract Configuration build(); - } - } - - /** The {@link Result} of dequeuing {@link RequestT}s. */ - static class Result implements POutput { - - static Result of(TupleTag requestsTag, PCollectionTuple pct) { - return new Result<>(requestsTag, pct); - } - - private final Pipeline pipeline; - private final TupleTag requestsTag; - private final PCollection requests; - private final PCollection failures; - - private Result(TupleTag requestsTag, PCollectionTuple pct) { - this.pipeline = pct.getPipeline(); - this.requestsTag = requestsTag; - this.requests = pct.get(requestsTag); - this.failures = pct.get(FAILURE_TAG); - } - - @Override - public Pipeline getPipeline() { - return pipeline; - } - - @Override - public Map, PValue> expand() { - return ImmutableMap.of( - requestsTag, requests, - FAILURE_TAG, failures); - } - - @Override - public void finishSpecifyingOutput( - String transformName, PInput input, PTransform transform) {} - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleEnqueue.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleEnqueue.java deleted file mode 100644 index 505ef86be48b3..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleEnqueue.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.value.AutoValue; -import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.TypeDescriptor; - -/** - * {@link ThrottleEnqueue} enqueues {@link RequestT} elements yielding an {@link ApiIOError} {@link - * PCollection} of any enqueue errors. - */ -class ThrottleEnqueue extends PTransform, PCollection> { - - @SuppressWarnings({"unused"}) - private final Configuration configuration; - - private ThrottleEnqueue(Configuration configuration) { - this.configuration = configuration; - } - - /** Configuration details for {@link ThrottleEnqueue}. */ - @AutoValue - abstract static class Configuration { - - static Builder builder() { - return new AutoValue_ThrottleEnqueue_Configuration.Builder<>(); - } - - abstract Builder toBuilder(); - - @AutoValue.Builder - abstract static class Builder { - - abstract Configuration build(); - } - } - - @Override - public PCollection expand(PCollection input) { - // TODO(damondouglas): expand in a future PR. - return input.getPipeline().apply(Create.empty(TypeDescriptor.of(ApiIOError.class))); - } -} diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleRefreshQuota.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleRefreshQuota.java deleted file mode 100644 index 57e57528db4bc..0000000000000 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleRefreshQuota.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * 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.io.requestresponse; - -import com.google.auto.value.AutoValue; -import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.TypeDescriptor; -import org.joda.time.Instant; - -/** - * {@link ThrottleRefreshQuota} refreshes a quota per {@link Instant} processing events emitting any - * errors into an {@link ApiIOError} {@link PCollection}. - */ -class ThrottleRefreshQuota extends PTransform, PCollection> { - - // TODO: remove suppress warnings after configuration utilized. - @SuppressWarnings({"unused"}) - private final Configuration configuration; - - private ThrottleRefreshQuota(Configuration configuration) { - this.configuration = configuration; - } - - @Override - public PCollection expand(PCollection input) { - // TODO(damondouglas): expand in a later PR. - return input.getPipeline().apply(Create.empty(TypeDescriptor.of(ApiIOError.class))); - } - - @AutoValue - abstract static class Configuration { - - @AutoValue.Builder - abstract static class Builder { - abstract Configuration build(); - } - } -} diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java deleted file mode 100644 index a53b74d86d7fb..0000000000000 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/DeterministicCoderCacheSerializerTest.java +++ /dev/null @@ -1,5 +0,0 @@ -package org.apache.beam.io.requestresponse; - -public class DeterministicCoderCacheSerializerTest { - -} \ No newline at end of file From 1c7b620b2fe3429eb0a848573c534b3625e13ca6 Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Fri, 10 Nov 2023 11:22:26 -0800 Subject: [PATCH 4/8] wip --- .../apache/beam/io/requestresponse/Cache.java | 124 +++++++++++------- .../apache/beam/io/requestresponse/Quota.java | 3 +- ...java => ThrottleWithExternalResource.java} | 34 ++--- .../beam/io/requestresponse/CacheTest.java | 101 ++++++++++++++ .../beam/io/requestresponse/CallTest.java | 42 +++++- .../ThrottleWithExternalResourceTest.java | 25 ++++ 6 files changed, 263 insertions(+), 66 deletions(-) rename sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/{Throttle.java => ThrottleWithExternalResource.java} (91%) create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java index 9760b91f13579..fd83a8d52bb4f 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java @@ -32,29 +32,35 @@ import org.joda.time.Duration; /** Transforms for reading and writing request/response associations to a cache. */ -final class Cache { +class Cache { /** * Read {@link RequestT} {@link ResponseT} associations from a cache. The {@link KV} value is null - * when no association exists. + * when no association exists. This method does not enforce {@link Coder#verifyDeterministic} and + * defers to the user to determine whether to enforce this given the cache implementation. */ static < @NonNull RequestT, @Nullable ResponseT, CallerSetupTeardownT extends - Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> & SetupTeardown> + @NonNull Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> + & SetupTeardown> PTransform< - PCollection<@NonNull RequestT>, - Call.Result>> + @NonNull PCollection<@NonNull RequestT>, + Call.@NonNull Result>> read( - CallerSetupTeardownT implementsCallerSetupTeardown, - Coder<@NonNull RequestT> requestTCoder, - Coder<@NonNull ResponseT> responseTCoder) { + @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@NonNull ResponseT> responseTCoder) { return Call.ofCallerAndSetupTeardown( implementsCallerSetupTeardown, KvCoder.of(requestTCoder, responseTCoder)); } - /** Write a {@link RequestT} {@link ResponseT} association to a cache. */ + /** + * Write a {@link RequestT} {@link ResponseT} association to a cache. This method does not enforce + * {@link Coder#verifyDeterministic} and defers to the user to determine whether to enforce this + * given the cache implementation. + */ static < @NonNull RequestT, @NonNull ResponseT, @@ -67,26 +73,34 @@ final class Cache { PCollection>, Call.Result>> write( - CallerSetupTeardownT implementsCallerSetupTeardown, - KvCoder<@NonNull RequestT, @NonNull ResponseT> kvCoder) { + @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, + @NonNull KvCoder<@NonNull RequestT, @NonNull ResponseT> kvCoder) { return Call.ofCallerAndSetupTeardown(implementsCallerSetupTeardown, kvCoder); } + static <@NonNull RequestT, ResponseT> UsingRedis usingRedis( + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull RedisClient client) + throws Coder.NonDeterministicException { + return new UsingRedis<>(requestTCoder, responseTCoder, client); + } + /** Provides cache read and write support using a {@link RedisClient}. */ - static class UsingRedis { - private final Coder<@NonNull RequestT> requestTCoder; - private final Coder<@NonNull ResponseT> responseTCoder; - private final RedisClient client; + static class UsingRedis<@NonNull RequestT, ResponseT> { + private final @NonNull Coder<@NonNull RequestT> requestTCoder; + private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull RedisClient client; /** - * Instantiates a {@link UsingRedis}. Throws a {@link Coder.NonDeterministicException} if either - * {@link RequestT} or {@link ResponseT} {@link Coder} fails to {@link - * Coder#verifyDeterministic}. + * Instantiates a {@link UsingRedis}. Given redis reads and writes using bytes, throws a {@link + * Coder.NonDeterministicException} if either {@link RequestT} or {@link ResponseT} {@link + * Coder} fails to {@link Coder#verifyDeterministic}. */ - UsingRedis( - Coder<@NonNull RequestT> requestTCoder, - Coder<@NonNull ResponseT> responseTCoder, - RedisClient client) + private UsingRedis( + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull RedisClient client) throws Coder.NonDeterministicException { this.client = client; requestTCoder.verifyDeterministic(); @@ -97,7 +111,7 @@ static class UsingRedis { /** Instantiates {@link Read}. */ Read<@NonNull RequestT, @Nullable ResponseT> read() { - return new Read<>(this); + return new Read<>(requestTCoder, responseTCoder, client); } /** @@ -105,7 +119,7 @@ static class UsingRedis { * RequestT} and {@link ResponseT} lasts in the cache. */ Write<@NonNull RequestT, @NonNull ResponseT> write(Duration expiry) { - return new Write<>(expiry, this); + return new Write<>(expiry, requestTCoder, responseTCoder, client); } /** Reads associated {@link RequestT} {@link ResponseT} using a {@link RedisClient}. */ @@ -113,10 +127,17 @@ static class Read<@NonNull RequestT, @Nullable ResponseT> implements Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>>, SetupTeardown { - private final UsingRedis spec; - - private Read(UsingRedis spec) { - this.spec = spec; + private final @NonNull Coder<@NonNull RequestT> requestTCoder; + private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull RedisClient client; + + private Read( + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull RedisClient client) { + this.requestTCoder = requestTCoder; + this.responseTCoder = responseTCoder; + this.client = client; } @Override @@ -124,15 +145,15 @@ private Read(UsingRedis spec) { throws UserCodeExecutionException { ByteArrayOutputStream baos = new ByteArrayOutputStream(); try { - spec.requestTCoder.encode(request, baos); + requestTCoder.encode(request, baos); byte[] encodedRequest = baos.toByteArray(); - byte[] encodedResponse = spec.client.getBytes(encodedRequest); + byte[] encodedResponse = client.getBytes(encodedRequest); if (encodedResponse == null) { return KV.of(request, null); } ResponseT response = checkStateNotNull( - spec.responseTCoder.decode(ByteSource.wrap(encodedResponse).openStream())); + responseTCoder.decode(ByteSource.wrap(encodedResponse).openStream())); return KV.of(request, response); } catch (IllegalStateException | IOException e) { throw new UserCodeExecutionException(e); @@ -141,52 +162,61 @@ private Read(UsingRedis spec) { @Override public void setup() throws UserCodeExecutionException { - spec.client.setup(); + client.setup(); } @Override public void teardown() throws UserCodeExecutionException { - spec.client.teardown(); + client.teardown(); } } } static class Write<@NonNull RequestT, @NonNull ResponseT> implements Caller< - KV<@NonNull RequestT, @NonNull ResponseT>, KV<@NonNull RequestT, @NonNull ResponseT>>, + @NonNull KV<@NonNull RequestT, @NonNull ResponseT>, + @NonNull KV<@NonNull RequestT, @NonNull ResponseT>>, SetupTeardown { - private final Duration expiry; - - private final UsingRedis spec; - - private Write(Duration expiry, UsingRedis spec) { + private final @NonNull Duration expiry; + private final @NonNull Coder<@NonNull RequestT> requestTCoder; + private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull RedisClient client; + + private Write( + @NonNull Duration expiry, + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull RedisClient client) { this.expiry = expiry; - this.spec = spec; + this.requestTCoder = requestTCoder; + this.responseTCoder = responseTCoder; + this.client = client; } @Override - public KV<@NonNull RequestT, @NonNull ResponseT> call( - KV<@NonNull RequestT, @NonNull ResponseT> request) throws UserCodeExecutionException { + public @NonNull KV<@NonNull RequestT, @NonNull ResponseT> call( + @NonNull KV<@NonNull RequestT, @NonNull ResponseT> request) + throws UserCodeExecutionException { ByteArrayOutputStream keyStream = new ByteArrayOutputStream(); ByteArrayOutputStream valueStream = new ByteArrayOutputStream(); try { - spec.requestTCoder.encode(request.getKey(), keyStream); - spec.responseTCoder.encode(request.getValue(), valueStream); + requestTCoder.encode(request.getKey(), keyStream); + responseTCoder.encode(request.getValue(), valueStream); } catch (IOException e) { throw new UserCodeExecutionException(e); } - spec.client.setex(keyStream.toByteArray(), valueStream.toByteArray(), expiry); + client.setex(keyStream.toByteArray(), valueStream.toByteArray(), expiry); return request; } @Override public void setup() throws UserCodeExecutionException { - spec.client.setup(); + client.setup(); } @Override public void teardown() throws UserCodeExecutionException { - spec.client.teardown(); + client.teardown(); } } } diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java index c4d2af113cebb..2e948e36e19d1 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java @@ -20,6 +20,7 @@ import java.io.Serializable; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; /** @@ -46,7 +47,7 @@ public long getNumRequests() { } @Override - public boolean equals(Object o) { + public boolean equals(@Nullable Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; Quota quota = (Quota) o; diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java similarity index 91% rename from sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java rename to sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java index eac9aacb4a3d9..d94b674a18a46 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Throttle.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java @@ -17,6 +17,9 @@ */ package org.apache.beam.io.requestresponse; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + import java.io.ByteArrayOutputStream; import java.io.IOException; import java.net.URI; @@ -82,26 +85,27 @@ class ThrottleWithExternalResource< private static final TupleTag FAILURE_TAG = new TupleTag() {}; private static final Duration THROTTLE_INTERVAL = Duration.standardSeconds(1L); - private final Quota quota; - private final String quotaIdentifier; - private final Coder coder; - private final ReporterT reporterT; - private final EnqueuerT enqueuerT; - private final DequeuerT dequeuerT; - private final RefresherT refresherT; + private final @NonNull Quota quota; + private final @NonNull String quotaIdentifier; + private final @NonNull Coder coder; + private final @NonNull ReporterT reporterT; + private final @NonNull EnqueuerT enqueuerT; + private final @NonNull DequeuerT dequeuerT; + private final @NonNull RefresherT refresherT; ThrottleWithExternalResource( - Quota quota, - String quotaIdentifier, - Coder coder, - ReporterT reporterT, - EnqueuerT enqueuerT, - DequeuerT dequeuerT, - RefresherT refresherT) + @NonNull Quota quota, + @NonNull String quotaIdentifier, + @NonNull Coder<@NonNull T> coder, + @NonNull ReporterT reporterT, + @NonNull EnqueuerT enqueuerT, + @NonNull DequeuerT dequeuerT, + @NonNull RefresherT refresherT) throws Coder.NonDeterministicException { this.quotaIdentifier = quotaIdentifier; this.reporterT = reporterT; coder.verifyDeterministic(); + checkArgument(!quotaIdentifier.isEmpty()); this.quota = quota; this.coder = coder; this.enqueuerT = enqueuerT; @@ -260,7 +264,7 @@ private RedisDequeur(URI uri, Coder coder, String key) { public @NonNull T call(@NonNull Instant request) throws UserCodeExecutionException { byte[] bytes = client.lpop(key); try { - return coder.decode(ByteSource.wrap(bytes).openStream()); + return checkStateNotNull(coder.decode(ByteSource.wrap(bytes).openStream())); } catch (IOException e) { throw new UserCodeExecutionException(e); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java new file mode 100644 index 0000000000000..a594573a8335a --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java @@ -0,0 +1,101 @@ +/* + * 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.io.requestresponse; + +import static org.junit.Assert.assertThrows; + +import java.net.URI; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.KV; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link Cache}. */ +@RunWith(JUnit4.class) +public class CacheTest { + @Rule public TestPipeline pipeline = TestPipeline.create(); + + @Test + public void givenNonDeterministicCoder_UsingRedis_throwsError() + throws Coder.NonDeterministicException { + assertThrows( + "Request Coder is non-deterministic", + Coder.NonDeterministicException.class, + () -> + Cache.usingRedis( + SerializableCoder.of(CallTest.Request.class), + StringUtf8Coder.of(), + new RedisClient(URI.create("redis://localhost:6379")))); + + assertThrows( + "Response Coder is non-deterministic", + Coder.NonDeterministicException.class, + () -> + Cache.usingRedis( + StringUtf8Coder.of(), + CallTest.RESPONSE_CODER, + new RedisClient(URI.create("redis://localhost:6379")))); + + // both Request and Response Coders are deterministic + Cache.usingRedis( + StringUtf8Coder.of(), + StringUtf8Coder.of(), + new RedisClient(URI.create("redis://localhost:6379"))); + } + + @Test + public void givenUsingRedis_pipelineConstructsDoesNotThrowError() { + } + + @Test + public void givenWrongRedisURI_throwsError() {} + + + private static class ValidCaller + implements Caller< + CallTest.@NonNull Request, + @NonNull KV>, + SetupTeardown { + private final @NonNull KV yieldValue; + + private ValidCaller( + @NonNull KV yieldValue) { + this.yieldValue = yieldValue; + } + + @Override + public void setup() throws UserCodeExecutionException {} + + @Override + public void teardown() throws UserCodeExecutionException {} + + @Override + public @NonNull KV call( + CallTest.@NonNull Request request) throws UserCodeExecutionException { + return yieldValue; + } + } +} diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java index 18574b00978dc..970955f1bc15a 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java @@ -22,9 +22,17 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.io.Serializable; import org.apache.beam.io.requestresponse.Call.Result; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.DelegateCoder; import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; @@ -47,7 +55,7 @@ public class CallTest { @Rule public TestPipeline pipeline = TestPipeline.create(); - private static final SerializableCoder<@NonNull Response> RESPONSE_CODER = + static final SerializableCoder<@NonNull Response> RESPONSE_CODER = SerializableCoder.of(Response.class); @Test @@ -275,7 +283,7 @@ public void teardown() throws UserCodeExecutionException {} private static class UnSerializable {} - private static class Request implements Serializable { + static class Request implements Serializable { final String id; @@ -305,7 +313,7 @@ public int hashCode() { } } - private static class Response implements Serializable { + static class Response implements Serializable { final String id; Response(String id) { @@ -490,4 +498,32 @@ private static void sleep(Duration timeout) { } catch (InterruptedException ignored) { } } + + private static class DeterministicRequestCoder extends CustomCoder { + private static final Coder ID_CODER = StringUtf8Coder.of(); + @Override + public void encode(Request value, OutputStream outStream) throws CoderException, IOException { + ID_CODER.encode(value.id, outStream); + } + + @Override + public Request decode(InputStream inStream) throws CoderException, IOException { + String id = ID_CODER.decode(inStream); + return new Request(id); + } + } + + private static class DeterministicResponseCoder extends CustomCoder { + private static final Coder ID_CODER = StringUtf8Coder.of(); + + @Override + public void encode(Response value, OutputStream outStream) throws CoderException, IOException { + + } + + @Override + public Response decode(InputStream inStream) throws CoderException, IOException { + return null; + } + } } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java new file mode 100644 index 0000000000000..4fd45abd27089 --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java @@ -0,0 +1,25 @@ +/* + * 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.io.requestresponse; + +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ThrottleWithExternalResource}. */ +@RunWith(JUnit4.class) +public class ThrottleWithExternalResourceTest {} From 7f22457ee9f976b48ad3be4fd59dad23c09b023a Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Sat, 11 Nov 2023 01:32:46 +0000 Subject: [PATCH 5/8] Implement Throttle and Cache --- .../configmap.yaml | 30 +++ .../deployment.yaml | 27 +++ .../kustomization.yaml | 34 ++++ .../apache/beam/io/requestresponse/Cache.java | 53 ++++-- .../apache/beam/io/requestresponse/Call.java | 34 ++-- .../apache/beam/io/requestresponse/Quota.java | 8 +- .../ThrottleWithExternalResource.java | 44 +++-- .../beam/io/requestresponse/CacheTest.java | 125 +++++++----- .../beam/io/requestresponse/CacheTestIT.java | 120 ++++++++++++ .../beam/io/requestresponse/CallTest.java | 94 ++++++--- .../io/requestresponse/EchoITOptions.java | 2 +- .../io/requestresponse/EchoRequestCoder.java | 44 +++++ .../ThrottleWithExternalResourceTest.java | 54 +++++- .../ThrottleWithExternalResourceTestIT.java | 180 ++++++++++++++++++ 14 files changed, 734 insertions(+), 115 deletions(-) create mode 100644 .test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/configmap.yaml create mode 100644 .test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/deployment.yaml create mode 100644 .test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/kustomization.yaml create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoRequestCoder.java create mode 100644 sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java diff --git a/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/configmap.yaml b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/configmap.yaml new file mode 100644 index 0000000000000..6a482b21b1657 --- /dev/null +++ b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/configmap.yaml @@ -0,0 +1,30 @@ +# 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. + +# Configures patch for ../base/configmap.yaml +# See https://kubectl.docs.kubernetes.io/references/kustomize/kustomization/patches/ + +- op: replace + path: /metadata/labels/quota-id + value: echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota +- op: replace + path: /data/QUOTA_ID + value: echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota +- op: replace + path: /data/QUOTA_SIZE + value: "10" +- op: replace + path: /data/QUOTA_REFRESH_INTERVAL + value: 1s diff --git a/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/deployment.yaml b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/deployment.yaml new file mode 100644 index 0000000000000..cff2f994cd63f --- /dev/null +++ b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/deployment.yaml @@ -0,0 +1,27 @@ +# 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. + +# Configures patch for ../base/deployment.yaml +# See https://kubectl.docs.kubernetes.io/references/kustomize/kustomization/patches/ + +- op: replace + path: /metadata/labels/quota-id + value: echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota +- op: replace + path: /spec/selector/matchLabels/quota-id + value: echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota +- op: replace + path: /spec/template/metadata/labels/quota-id + value: echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota diff --git a/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/kustomization.yaml b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/kustomization.yaml new file mode 100644 index 0000000000000..d10598be51f47 --- /dev/null +++ b/.test-infra/mock-apis/infrastructure/kubernetes/refresher/overlays/echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota/kustomization.yaml @@ -0,0 +1,34 @@ +# 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. + +# Configures the overlay for .test-infra/mock-apis/infrastructure/kubernetes/refresher/base +# Using the Quota Id: +# echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota + +resources: +- ../../base + +nameSuffix: -throttle-with-external-resource-test-10-per-1s + +patches: +- path: configmap.yaml + target: + kind: ConfigMap + name: refresher + +- path: deployment.yaml + target: + kind: Deployment + name: refresher diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java index fd83a8d52bb4f..1eb09a9386535 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java @@ -22,6 +22,7 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.values.KV; @@ -51,11 +52,26 @@ class Cache { read( @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@NonNull ResponseT> responseTCoder) { + @NonNull Coder<@Nullable ResponseT> responseTCoder) { return Call.ofCallerAndSetupTeardown( implementsCallerSetupTeardown, KvCoder.of(requestTCoder, responseTCoder)); } + static <@NonNull RequestT, @Nullable ResponseT> + PTransform< + @NonNull PCollection<@NonNull RequestT>, + Call.@NonNull Result>> + readUsingRedis( + @NonNull RedisClient client, + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder) + throws NonDeterministicException { + return read( + new UsingRedis<>(requestTCoder, responseTCoder, client).read(), + requestTCoder, + responseTCoder); + } + /** * Write a {@link RequestT} {@link ResponseT} association to a cache. This method does not enforce * {@link Coder#verifyDeterministic} and defers to the user to determine whether to enforce this @@ -70,17 +86,32 @@ class Cache { KV<@NonNull RequestT, @NonNull ResponseT>> & SetupTeardown> PTransform< - PCollection>, - Call.Result>> + @NonNull PCollection>, + Call.@NonNull Result>> write( @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, @NonNull KvCoder<@NonNull RequestT, @NonNull ResponseT> kvCoder) { return Call.ofCallerAndSetupTeardown(implementsCallerSetupTeardown, kvCoder); } + static <@NonNull RequestT, @NonNull ResponseT> + PTransform< + @NonNull PCollection>, + Call.@NonNull Result>> + writeUsingRedis( + @NonNull Duration expiry, + @NonNull RedisClient client, + @NonNull Coder<@NonNull RequestT> requestTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder) + throws NonDeterministicException { + return write( + new UsingRedis<>(requestTCoder, responseTCoder, client).write(expiry), + KvCoder.of(requestTCoder, responseTCoder)); + } + static <@NonNull RequestT, ResponseT> UsingRedis usingRedis( @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder, @NonNull RedisClient client) throws Coder.NonDeterministicException { return new UsingRedis<>(requestTCoder, responseTCoder, client); @@ -89,7 +120,7 @@ class Cache { /** Provides cache read and write support using a {@link RedisClient}. */ static class UsingRedis<@NonNull RequestT, ResponseT> { private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull Coder<@Nullable ResponseT> responseTCoder; private final @NonNull RedisClient client; /** @@ -99,7 +130,7 @@ static class UsingRedis<@NonNull RequestT, ResponseT> { */ private UsingRedis( @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder, @NonNull RedisClient client) throws Coder.NonDeterministicException { this.client = client; @@ -118,7 +149,7 @@ private UsingRedis( * Instantiates {@link Write}. The {@link Duration} determines how long the associated {@link * RequestT} and {@link ResponseT} lasts in the cache. */ - Write<@NonNull RequestT, @NonNull ResponseT> write(Duration expiry) { + Write<@NonNull RequestT, @NonNull ResponseT> write(@NonNull Duration expiry) { return new Write<>(expiry, requestTCoder, responseTCoder, client); } @@ -128,12 +159,12 @@ static class Read<@NonNull RequestT, @Nullable ResponseT> SetupTeardown { private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull Coder<@Nullable ResponseT> responseTCoder; private final @NonNull RedisClient client; private Read( @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder, @NonNull RedisClient client) { this.requestTCoder = requestTCoder; this.responseTCoder = responseTCoder; @@ -179,13 +210,13 @@ static class Write<@NonNull RequestT, @NonNull ResponseT> SetupTeardown { private final @NonNull Duration expiry; private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@NonNull ResponseT> responseTCoder; + private final @NonNull Coder<@Nullable ResponseT> responseTCoder; private final @NonNull RedisClient client; private Write( @NonNull Duration expiry, @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@NonNull ResponseT> responseTCoder, + @NonNull Coder<@Nullable ResponseT> responseTCoder, @NonNull RedisClient client) { this.expiry = expiry; this.requestTCoder = requestTCoder; diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java index 52181af534ed3..d6d0881c11f9d 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java @@ -100,7 +100,7 @@ Call ofCallerAndSetupTeardown( .build()); } - private static final TupleTag FAILURE_TAG = new TupleTag() {}; + private final TupleTag failureTag = new TupleTag() {}; private final Configuration configuration; @@ -133,22 +133,26 @@ Call withTimeout(Duration timeout) { PCollectionTuple pct = input.apply( CallFn.class.getSimpleName(), - ParDo.of(new CallFn<>(responseTag, configuration)) - .withOutputTags(responseTag, TupleTagList.of(FAILURE_TAG))); + ParDo.of(new CallFn<>(responseTag, failureTag, configuration)) + .withOutputTags(responseTag, TupleTagList.of(failureTag))); - return Result.of(configuration.getResponseCoder(), responseTag, pct); + return Result.of(configuration.getResponseCoder(), responseTag, failureTag, pct); } private static class CallFn extends DoFn { private final TupleTag responseTag; + private final TupleTag failureTag; private final CallerWithTimeout caller; private final SetupTeardownWithTimeout setupTeardown; private transient @MonotonicNonNull ExecutorService executor; private CallFn( - TupleTag responseTag, Configuration configuration) { + TupleTag responseTag, + TupleTag failureTag, + Configuration configuration) { this.responseTag = responseTag; + this.failureTag = failureTag; this.caller = new CallerWithTimeout<>(configuration.getTimeout(), configuration.getCaller()); this.setupTeardown = new SetupTeardownWithTimeout( @@ -194,7 +198,7 @@ public void process(@Element @NonNull RequestT request, MultiOutputReceiver rece ResponseT response = this.caller.call(request); receiver.get(responseTag).output(response); } catch (UserCodeExecutionException e) { - receiver.get(FAILURE_TAG).output(ApiIOError.of(e, request)); + receiver.get(failureTag).output(ApiIOError.of(e, request)); } } } @@ -269,21 +273,29 @@ final Configuration build() { static class Result implements POutput { static Result of( - Coder responseTCoder, TupleTag responseTag, PCollectionTuple pct) { - return new Result<>(responseTCoder, responseTag, pct); + Coder responseTCoder, + TupleTag responseTag, + TupleTag failureTag, + PCollectionTuple pct) { + return new Result<>(responseTCoder, responseTag, pct, failureTag); } private final Pipeline pipeline; private final TupleTag responseTag; + private final TupleTag failureTag; private final PCollection responses; private final PCollection failures; private Result( - Coder responseTCoder, TupleTag responseTag, PCollectionTuple pct) { + Coder responseTCoder, + TupleTag responseTag, + PCollectionTuple pct, + TupleTag failureTag) { this.pipeline = pct.getPipeline(); this.responseTag = responseTag; + this.failureTag = failureTag; this.responses = pct.get(responseTag).setCoder(responseTCoder); - this.failures = pct.get(FAILURE_TAG); + this.failures = pct.get(this.failureTag); } public PCollection getResponses() { @@ -303,7 +315,7 @@ public PCollection getFailures() { public @NonNull Map, PValue> expand() { return ImmutableMap.of( responseTag, responses, - FAILURE_TAG, failures); + failureTag, failures); } @Override diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java index 2e948e36e19d1..d2e538cf7cf39 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Quota.java @@ -48,8 +48,12 @@ public long getNumRequests() { @Override public boolean equals(@Nullable Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } Quota quota = (Quota) o; return numRequests == quota.numRequests && Objects.equal(interval, quota.interval); } diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java index d94b674a18a46..25c94c4c31369 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java @@ -49,7 +49,7 @@ class ThrottleWithExternalResource< EnqueuerT extends Caller<@NonNull T, Void> & SetupTeardown, DequeuerT extends Caller<@NonNull Instant, @NonNull T> & SetupTeardown, RefresherT extends Caller<@NonNull Instant, Void> & SetupTeardown> - extends PTransform, Call.Result<@NonNull T>> { + extends PTransform<@NonNull PCollection<@NonNull T>, Call.@NonNull Result<@NonNull T>> { /** Instantiate a {@link ThrottleWithExternalResource} using a {@link RedisClient}. */ static <@NonNull T> @@ -81,13 +81,11 @@ class ThrottleWithExternalResource< new RedisRefresher(uri, quota, quotaIdentifier)); } - private final TupleTag<@NonNull T> outputTag = new TupleTag<@NonNull T>() {}; - private static final TupleTag FAILURE_TAG = new TupleTag() {}; private static final Duration THROTTLE_INTERVAL = Duration.standardSeconds(1L); private final @NonNull Quota quota; private final @NonNull String quotaIdentifier; - private final @NonNull Coder coder; + private final @NonNull Coder<@NonNull T> coder; private final @NonNull ReporterT reporterT; private final @NonNull EnqueuerT enqueuerT; private final @NonNull DequeuerT dequeuerT; @@ -114,7 +112,7 @@ class ThrottleWithExternalResource< } @Override - public Call.Result<@NonNull T> expand(PCollection<@NonNull T> input) { + public Call.@NonNull Result<@NonNull T> expand(PCollection<@NonNull T> input) { Pipeline pipeline = input.getPipeline(); // Refresh known quota to control the throttle rate. @@ -126,25 +124,34 @@ class ThrottleWithExternalResource< // Enqueue T elements. Call.Result enqueuResult = input.apply("enqueue", getEnqueuer()); + TupleTag<@NonNull T> outputTag = new TupleTag<@NonNull T>() {}; + TupleTag<@NonNull ApiIOError> failureTag = new TupleTag<@NonNull ApiIOError>() {}; + // Perform Throttle. PCollectionTuple pct = pipeline .apply("throttle/impulse", PeriodicImpulse.create().withInterval(THROTTLE_INTERVAL)) .apply( "throttle/fn", - ParDo.of(new ThrottleFn(quotaIdentifier, dequeuerT, reporterT, outputTag)) - .withOutputTags(outputTag, TupleTagList.of(FAILURE_TAG))); + ParDo.of( + new ThrottleFn( + quotaIdentifier, dequeuerT, reporterT, outputTag, failureTag)) + .withOutputTags(outputTag, TupleTagList.of(failureTag))); PCollection errors = PCollectionList.of(refreshResult.getFailures()) .and(enqueuResult.getFailures()) - .and(pct.get(FAILURE_TAG)) + .and(pct.get(failureTag)) .apply("errors/flatten", Flatten.pCollections()); - return Call.Result.of( + TupleTag<@NonNull T> resultOutputTag = new TupleTag<@NonNull T>() {}; + TupleTag<@NonNull ApiIOError> resultFailureTag = new TupleTag<@NonNull ApiIOError>() {}; + + return Call.Result.<@NonNull T>of( coder, - outputTag, - PCollectionTuple.of(outputTag, pct.get(outputTag)).and(FAILURE_TAG, errors)); + resultOutputTag, + resultFailureTag, + PCollectionTuple.of(resultOutputTag, pct.get(outputTag)).and(resultFailureTag, errors)); } private Call<@NonNull Instant, Void> getRefresher() { @@ -160,16 +167,19 @@ private class ThrottleFn extends DoFn<@NonNull Instant, @NonNull T> { private final DequeuerT dequeuerT; private final ReporterT reporterT; private final TupleTag<@NonNull T> outputTag; + private final TupleTag<@NonNull ApiIOError> failureTag; private ThrottleFn( String quotaIdentifier, DequeuerT dequeuerT, ReporterT reporterT, - TupleTag<@NonNull T> outputTag) { + TupleTag<@NonNull T> outputTag, + TupleTag<@NonNull ApiIOError> failureTag) { this.quotaIdentifier = quotaIdentifier; this.dequeuerT = dequeuerT; this.reporterT = reporterT; this.outputTag = outputTag; + this.failureTag = failureTag; } @ProcessElement @@ -186,7 +196,7 @@ public void process(@Element @NonNull Instant instant, MultiOutputReceiver recei receiver.get(outputTag).output(element); } catch (UserCodeExecutionException e) { receiver - .get(FAILURE_TAG) + .get(failureTag) .output( ApiIOError.builder() .setRequestAsJsonString(String.format("{\"request\": \"%s\"}", instant)) @@ -227,9 +237,9 @@ private RedisReporter(URI uri) { private static class RedisEnqueuer<@NonNull T> extends RedisSetupTeardown implements Caller<@NonNull T, Void> { private final String key; - private final Coder coder; + private final Coder<@NonNull T> coder; - private RedisEnqueuer(URI uri, String key, Coder coder) { + private RedisEnqueuer(URI uri, String key, Coder<@NonNull T> coder) { super(new RedisClient(uri)); this.key = key; this.coder = coder; @@ -251,10 +261,10 @@ public Void call(@NonNull T request) throws UserCodeExecutionException { private static class RedisDequeur<@NonNull T> extends RedisSetupTeardown implements Caller<@NonNull Instant, @NonNull T> { - private final Coder coder; + private final Coder<@NonNull T> coder; private final String key; - private RedisDequeur(URI uri, Coder coder, String key) { + private RedisDequeur(URI uri, Coder<@NonNull T> coder, String key) { super(new RedisClient(uri)); this.coder = coder; this.key = key; diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java index a594573a8335a..fcb7862e991ee 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTest.java @@ -17,17 +17,21 @@ */ package org.apache.beam.io.requestresponse; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; import static org.junit.Assert.assertThrows; import java.net.URI; +import org.apache.beam.io.requestresponse.CallTest.Request; +import org.apache.beam.io.requestresponse.CallTest.Response; import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.SerializableCoder; -import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.values.KV; -import org.checkerframework.checker.nullness.qual.NonNull; -import org.checkerframework.checker.nullness.qual.Nullable; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; +import org.joda.time.Duration; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; @@ -39,63 +43,90 @@ public class CacheTest { @Rule public TestPipeline pipeline = TestPipeline.create(); @Test - public void givenNonDeterministicCoder_UsingRedis_throwsError() + public void givenNonDeterministicCoder_readUsingRedis_throwsError() throws Coder.NonDeterministicException { + URI uri = URI.create("redis://localhost:6379"); assertThrows( - "Request Coder is non-deterministic", - Coder.NonDeterministicException.class, + NonDeterministicException.class, () -> - Cache.usingRedis( - SerializableCoder.of(CallTest.Request.class), - StringUtf8Coder.of(), - new RedisClient(URI.create("redis://localhost:6379")))); + Cache.readUsingRedis( + new RedisClient(uri), + CallTest.NON_DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); assertThrows( - "Response Coder is non-deterministic", - Coder.NonDeterministicException.class, + NonDeterministicException.class, () -> - Cache.usingRedis( - StringUtf8Coder.of(), - CallTest.RESPONSE_CODER, - new RedisClient(URI.create("redis://localhost:6379")))); + Cache.readUsingRedis( + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.NON_DETERMINISTIC_RESPONSE_CODER)); - // both Request and Response Coders are deterministic - Cache.usingRedis( - StringUtf8Coder.of(), - StringUtf8Coder.of(), - new RedisClient(URI.create("redis://localhost:6379"))); + Cache.readUsingRedis( + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER); } @Test - public void givenUsingRedis_pipelineConstructsDoesNotThrowError() { - } - - @Test - public void givenWrongRedisURI_throwsError() {} - + public void givenNonDeterministicCoder_writeUsingRedis_throwsError() + throws Coder.NonDeterministicException { + URI uri = URI.create("redis://localhost:6379"); + Duration expiry = Duration.standardSeconds(1L); + assertThrows( + NonDeterministicException.class, + () -> + Cache.writeUsingRedis( + expiry, + new RedisClient(uri), + CallTest.NON_DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); - private static class ValidCaller - implements Caller< - CallTest.@NonNull Request, - @NonNull KV>, - SetupTeardown { - private final @NonNull KV yieldValue; + assertThrows( + NonDeterministicException.class, + () -> + Cache.writeUsingRedis( + expiry, + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.NON_DETERMINISTIC_RESPONSE_CODER)); - private ValidCaller( - @NonNull KV yieldValue) { - this.yieldValue = yieldValue; - } + Cache.writeUsingRedis( + expiry, + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER); + } - @Override - public void setup() throws UserCodeExecutionException {} + @Test + public void givenWrongRedisURI_throwsError() throws NonDeterministicException { + URI uri = URI.create("redis://1.2.3.4:6379"); + Duration expiry = Duration.standardSeconds(1L); + PCollection requests = + pipeline + .apply("create requests", Create.of(new Request(""))) + .setCoder(CallTest.DETERMINISTIC_REQUEST_CODER); + requests.apply( + "readUsingRedis", + Cache.readUsingRedis( + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); - @Override - public void teardown() throws UserCodeExecutionException {} + PCollection> kvs = + pipeline.apply("create kvs", Create.of(KV.of(new Request(""), new Response("")))); + kvs.apply( + "writeUsingRedis", + Cache.writeUsingRedis( + expiry, + new RedisClient(uri), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); - @Override - public @NonNull KV call( - CallTest.@NonNull Request request) throws UserCodeExecutionException { - return yieldValue; - } + UncheckedExecutionException error = + assertThrows(UncheckedExecutionException.class, pipeline::run); + assertThat( + error.getCause().getMessage(), + containsString("Failed to connect to host: redis://1.2.3.4:6379")); } } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java new file mode 100644 index 0000000000000..fed09a2a19041 --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java @@ -0,0 +1,120 @@ +/* + * 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.io.requestresponse; + +import java.net.URI; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.io.requestresponse.CallTest.Request; +import org.apache.beam.io.requestresponse.CallTest.Response; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.joda.time.Duration; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.utility.DockerImageName; + +/** Integration tests for {@link Cache}. */ +@RunWith(JUnit4.class) +public class CacheTestIT { + @Rule public TestPipeline writePipeline = TestPipeline.create(); + + @Rule public TestPipeline readPipeline = TestPipeline.create(); + + private static final String CONTAINER_IMAGE_NAME = "redis:5.0.3-alpine"; + private static final Integer PORT = 6379; + + @Rule + public GenericContainer redis = + new GenericContainer<>(DockerImageName.parse(CONTAINER_IMAGE_NAME)).withExposedPorts(PORT); + + @Rule + public RedisExternalResourcesRule externalClients = + new RedisExternalResourcesRule( + () -> { + redis.start(); + return URI.create( + String.format("redis://%s:%d", redis.getHost(), redis.getFirstMappedPort())); + }); + + @Test + public void givenRequestResponsesCached_writeThenReadYieldsMatches() + throws NonDeterministicException { + List> toWrite = + ImmutableList.of( + KV.of(new Request("a"), new Response("a")), + KV.of(new Request("b"), new Response("b")), + KV.of(new Request("c"), new Response("c"))); + List toRead = ImmutableList.of(new Request("a"), new Request("b"), new Request("c")); + writeThenReadThenPAssert(toWrite, toRead, toWrite); + } + + @Test + public void givenNoMatchingRequestResponsePairs_yieldsKVsWithNullValues() + throws NonDeterministicException { + List> toWrite = + ImmutableList.of( + KV.of(new Request("a"), new Response("a")), + KV.of(new Request("b"), new Response("b")), + KV.of(new Request("c"), new Response("c"))); + List toRead = ImmutableList.of(new Request("d"), new Request("e"), new Request("f")); + List> expected = + toRead.stream() + .>map(request -> KV.of(request, null)) + .collect(Collectors.toList()); + writeThenReadThenPAssert(toWrite, toRead, expected); + } + + private void writeThenReadThenPAssert( + List> toWrite, + List toRead, + List> expected) + throws NonDeterministicException { + PCollection> toWritePCol = writePipeline.apply(Create.of(toWrite)); + toWritePCol.apply( + Cache.writeUsingRedis( + Duration.standardHours(1L), + externalClients.getActualClient(), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); + + PCollection requests = + readPipeline.apply(Create.of(toRead)).setCoder(CallTest.DETERMINISTIC_REQUEST_CODER); + + Call.Result> gotKVsResult = + requests.apply( + Cache.readUsingRedis( + externalClients.getActualClient(), + CallTest.DETERMINISTIC_REQUEST_CODER, + CallTest.DETERMINISTIC_RESPONSE_CODER)); + + PAssert.that(gotKVsResult.getFailures()).empty(); + PAssert.that(gotKVsResult.getResponses()).containsInAnyOrder(expected); + + writePipeline.run().waitUntilFinish(); + readPipeline.run(); + } +} diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java index 970955f1bc15a..1314c1b772903 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java @@ -30,7 +30,7 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.CustomCoder; -import org.apache.beam.sdk.coders.DelegateCoder; +import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.testing.PAssert; @@ -43,7 +43,11 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; +import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.checkerframework.checker.nullness.qual.UnknownKeyFor; +import org.jetbrains.annotations.NotNull; import org.joda.time.Duration; import org.junit.Rule; import org.junit.Test; @@ -55,13 +59,23 @@ public class CallTest { @Rule public TestPipeline pipeline = TestPipeline.create(); - static final SerializableCoder<@NonNull Response> RESPONSE_CODER = + static final SerializableCoder<@NonNull Request> NON_DETERMINISTIC_REQUEST_CODER = + SerializableCoder.of(Request.class); + + static final Coder<@NonNull Request> DETERMINISTIC_REQUEST_CODER = + new DeterministicRequestCoder(); + + static final SerializableCoder<@NonNull Response> NON_DETERMINISTIC_RESPONSE_CODER = SerializableCoder.of(Response.class); + static final Coder<@NonNull Response> DETERMINISTIC_RESPONSE_CODER = + new DeterministicResponseCoder(); + @Test public void givenCallerNotSerializable_throwsError() { assertThrows( - IllegalArgumentException.class, () -> Call.of(new UnSerializableCaller(), RESPONSE_CODER)); + IllegalArgumentException.class, + () -> Call.of(new UnSerializableCaller(), NON_DETERMINISTIC_RESPONSE_CODER)); } @Test @@ -70,7 +84,7 @@ public void givenSetupTeardownNotSerializable_throwsError() { IllegalArgumentException.class, () -> Call.ofCallerAndSetupTeardown( - new UnSerializableCallerWithSetupTeardown(), RESPONSE_CODER)); + new UnSerializableCallerWithSetupTeardown(), NON_DETERMINISTIC_RESPONSE_CODER)); } @Test @@ -78,7 +92,10 @@ public void givenCallerThrowsUserCodeExecutionException_emitsIntoFailurePCollect Result result = pipeline .apply(Create.of(new Request("a"))) - .apply(Call.of(new CallerThrowsUserCodeExecutionException(), RESPONSE_CODER)); + .apply( + Call.of( + new CallerThrowsUserCodeExecutionException(), + NON_DETERMINISTIC_RESPONSE_CODER)); PCollection failures = result.getFailures(); PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class)) @@ -95,7 +112,7 @@ public void givenCallerThrowsQuotaException_emitsIntoFailurePCollection() { Result result = pipeline .apply(Create.of(new Request("a"))) - .apply(Call.of(new CallerInvokesQuotaException(), RESPONSE_CODER)); + .apply(Call.of(new CallerInvokesQuotaException(), NON_DETERMINISTIC_RESPONSE_CODER)); PCollection failures = result.getFailures(); PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class)) @@ -113,7 +130,9 @@ public void givenCallerTimeout_emitsFailurePCollection() { Result result = pipeline .apply(Create.of(new Request("a"))) - .apply(Call.of(new CallerExceedsTimeout(timeout), RESPONSE_CODER).withTimeout(timeout)); + .apply( + Call.of(new CallerExceedsTimeout(timeout), NON_DETERMINISTIC_RESPONSE_CODER) + .withTimeout(timeout)); PCollection failures = result.getFailures(); PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class)) @@ -130,7 +149,7 @@ public void givenCallerThrowsTimeoutException_emitsFailurePCollection() { Result result = pipeline .apply(Create.of(new Request("a"))) - .apply(Call.of(new CallerThrowsTimeout(), RESPONSE_CODER)); + .apply(Call.of(new CallerThrowsTimeout(), NON_DETERMINISTIC_RESPONSE_CODER)); PCollection failures = result.getFailures(); PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class)) @@ -147,7 +166,7 @@ public void givenSetupThrowsUserCodeExecutionException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new SetupThrowsUserCodeExecutionException())); assertPipelineThrows(UserCodeExecutionException.class, pipeline); @@ -158,7 +177,7 @@ public void givenSetupThrowsQuotaException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new SetupThrowsUserCodeQuotaException())); assertPipelineThrows(UserCodeQuotaException.class, pipeline); @@ -171,7 +190,7 @@ public void givenSetupTimeout_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new SetupExceedsTimeout(timeout)) .withTimeout(timeout)); @@ -183,7 +202,7 @@ public void givenSetupThrowsTimeoutException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new SetupThrowsUserCodeTimeoutException())); assertPipelineThrows(UserCodeTimeoutException.class, pipeline); @@ -194,7 +213,7 @@ public void givenTeardownThrowsUserCodeExecutionException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new TeardownThrowsUserCodeExecutionException())); // Exceptions thrown during teardown do not populate with the cause @@ -206,7 +225,7 @@ public void givenTeardownThrowsQuotaException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new TeardownThrowsUserCodeQuotaException())); // Exceptions thrown during teardown do not populate with the cause @@ -219,7 +238,7 @@ public void givenTeardownTimeout_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withTimeout(timeout) .withSetupTeardown(new TeardownExceedsTimeout(timeout))); @@ -232,7 +251,7 @@ public void givenTeardownThrowsTimeoutException_throwsError() { pipeline .apply(Create.of(new Request(""))) .apply( - Call.of(new ValidCaller(), RESPONSE_CODER) + Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER) .withSetupTeardown(new TeardownThrowsUserCodeTimeoutException())); // Exceptions thrown during teardown do not populate with the cause @@ -244,7 +263,7 @@ public void givenValidCaller_emitValidResponse() { Result result = pipeline .apply(Create.of(new Request("a"))) - .apply(Call.of(new ValidCaller(), RESPONSE_CODER)); + .apply(Call.of(new ValidCaller(), NON_DETERMINISTIC_RESPONSE_CODER)); PAssert.thatSingleton(result.getFailures().apply(Count.globally())).isEqualTo(0L); PAssert.that(result.getResponses()).containsInAnyOrder(new Response("a")); @@ -499,31 +518,56 @@ private static void sleep(Duration timeout) { } } - private static class DeterministicRequestCoder extends CustomCoder { + private static class DeterministicRequestCoder extends CustomCoder<@NonNull Request> { private static final Coder ID_CODER = StringUtf8Coder.of(); + @Override - public void encode(Request value, OutputStream outStream) throws CoderException, IOException { - ID_CODER.encode(value.id, outStream); + public void encode(Request value, @NotNull OutputStream outStream) + throws CoderException, IOException { + ID_CODER.encode(checkStateNotNull(value).id, outStream); } @Override - public Request decode(InputStream inStream) throws CoderException, IOException { + public @NonNull Request decode(@NotNull InputStream inStream) + throws CoderException, IOException { String id = ID_CODER.decode(inStream); return new Request(id); } + + @Override + public void verifyDeterministic() + throws @UnknownKeyFor @NonNull @Initialized NonDeterministicException { + ID_CODER.verifyDeterministic(); + } } private static class DeterministicResponseCoder extends CustomCoder { - private static final Coder ID_CODER = StringUtf8Coder.of(); + private static final NullableCoder ID_CODER = NullableCoder.of(StringUtf8Coder.of()); @Override - public void encode(Response value, OutputStream outStream) throws CoderException, IOException { + public void encode(@Nullable Response value, @NotNull OutputStream outStream) + throws CoderException, IOException { + if (value == null) { + ID_CODER.encode(null, outStream); + return; + } + ID_CODER.encode(checkStateNotNull(value).id, outStream); + } + @Override + public Response decode(@NotNull InputStream inStream) throws CoderException, IOException { + try { + String id = ID_CODER.decode(inStream); + return new Response(id); + } catch (CoderException ignored) { + return null; + } } @Override - public Response decode(InputStream inStream) throws CoderException, IOException { - return null; + public void verifyDeterministic() + throws @UnknownKeyFor @NonNull @Initialized NonDeterministicException { + ID_CODER.verifyDeterministic(); } } } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java index a32f7a78e8265..0219d5d4f6341 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java @@ -31,7 +31,7 @@ * * 
  *   ./gradlew :sdks:java:io:rrio:integrationTest -DintegrationTestPipelineOptions='[
- *      "--grpcEndpointAddress=",
+ *      "--gRPCEndpointAddress=",
  *      "--httpEndpointAddress="
  *   ]'
  *
diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoRequestCoder.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoRequestCoder.java new file mode 100644 index 0000000000000..75cd49904cff8 --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoRequestCoder.java @@ -0,0 +1,44 @@ +/* + * 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.io.requestresponse; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoRequest; +import org.checkerframework.checker.nullness.qual.NonNull; + +class EchoRequestCoder extends CustomCoder<@NonNull EchoRequest> { + + @Override + public void encode(@NonNull EchoRequest value, @NonNull OutputStream outStream) + throws CoderException, IOException { + value.writeTo(outStream); + } + + @Override + public @NonNull EchoRequest decode(@NonNull InputStream inStream) + throws CoderException, IOException { + return EchoRequest.parseFrom(inStream); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException {} +} diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java index 4fd45abd27089..591ba923201fc 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTest.java @@ -17,9 +17,61 @@ */ package org.apache.beam.io.requestresponse; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertThrows; + +import java.net.URI; +import org.apache.beam.io.requestresponse.CallTest.Request; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; +import org.joda.time.Duration; +import org.junit.Rule; +import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; /** Tests for {@link ThrottleWithExternalResource}. */ @RunWith(JUnit4.class) -public class ThrottleWithExternalResourceTest {} +public class ThrottleWithExternalResourceTest { + @Rule public TestPipeline pipeline = TestPipeline.create(); + + @Test + public void givenNonDeterministicCoder_usingRedis_throwsError() throws NonDeterministicException { + URI uri = URI.create("redis://localhost:6379"); + String quotaIdentifier = "quota"; + String queueKey = "queue"; + Quota quota = new Quota(10L, Duration.standardSeconds(1L)); + + assertThrows( + NonDeterministicException.class, + () -> + ThrottleWithExternalResource.usingRedis( + uri, quotaIdentifier, queueKey, quota, CallTest.NON_DETERMINISTIC_REQUEST_CODER)); + + ThrottleWithExternalResource.usingRedis( + uri, quotaIdentifier, queueKey, quota, CallTest.DETERMINISTIC_REQUEST_CODER); + } + + @Test + public void givenWrongRedisURI_throwsError() throws NonDeterministicException { + URI uri = URI.create("redis://1.2.3.4:6379"); + String quotaIdentifier = "quota"; + String queueKey = "queue"; + Quota quota = new Quota(10L, Duration.standardSeconds(1L)); + PCollection requests = + pipeline.apply(Create.of(new Request(""))).setCoder(CallTest.DETERMINISTIC_REQUEST_CODER); + requests.apply( + ThrottleWithExternalResource.usingRedis( + uri, quotaIdentifier, queueKey, quota, requests.getCoder())); + + UncheckedExecutionException error = + assertThrows(UncheckedExecutionException.class, pipeline::run); + assertThat( + error.getCause().getMessage(), + containsString("Failed to connect to host: redis://1.2.3.4:6379")); + } +} diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java new file mode 100644 index 0000000000000..007e418f21393 --- /dev/null +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java @@ -0,0 +1,180 @@ +/* + * 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.io.requestresponse; + +import static org.apache.beam.sdk.io.common.IOITHelper.readIOTestPipelineOptions; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.sdk.values.TypeDescriptors.strings; + +import com.google.protobuf.ByteString; +import java.net.URI; +import java.util.UUID; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.NonDeterministicException; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Distinct; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.PeriodicImpulse; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoRequest; +import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoResponse; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.joda.time.Duration; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.utility.DockerImageName; + +/** + * Integration tests for {@link ThrottleWithExternalResource}. See {@link EchoITOptions} for details + * on the required parameters and how to provide these for running integration tests. + */ +public class ThrottleWithExternalResourceTestIT { + + @Rule public TestPipeline pipeline = TestPipeline.create(); + + private static final String QUOTA_ID = "echo-ThrottleWithExternalResourceTestIT-10-per-1s-quota"; + private static final Quota QUOTA = new Quota(1L, Duration.standardSeconds(1L)); + private static final ByteString PAYLOAD = ByteString.copyFromUtf8("payload"); + private static @MonotonicNonNull EchoITOptions options; + private static @MonotonicNonNull EchoGRPCCallerWithSetupTeardown client; + private static final String CONTAINER_IMAGE_NAME = "redis:5.0.3-alpine"; + private static final Integer PORT = 6379; + private static final EchoRequestCoder REQUEST_CODER = new EchoRequestCoder(); + private static final Coder RESPONSE_CODER = + SerializableCoder.of(TypeDescriptor.of(EchoResponse.class)); + + @Rule + public GenericContainer redis = + new GenericContainer<>(DockerImageName.parse(CONTAINER_IMAGE_NAME)).withExposedPorts(PORT); + + @BeforeClass + public static void setUp() throws UserCodeExecutionException { + options = readIOTestPipelineOptions(EchoITOptions.class); + if (options.getgRPCEndpointAddress().isEmpty()) { + throw new RuntimeException( + "--gRPCEndpointAddress is missing. See " + EchoITOptions.class + "for details."); + } + client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getgRPCEndpointAddress())); + checkStateNotNull(client).setup(); + + try { + client.call(createRequest()); + } catch (UserCodeExecutionException e) { + if (e instanceof UserCodeQuotaException) { + throw new RuntimeException( + String.format( + "The quota: %s is set to refresh on an interval. Unless there are failures in this test, wait for a few seconds before running the test again.", + QUOTA_ID), + e); + } + throw e; + } + } + + @AfterClass + public static void tearDown() throws UserCodeExecutionException { + checkStateNotNull(client).teardown(); + } + + @Test + public void givenThrottleUsingRedis_preventsQuotaErrors() throws NonDeterministicException { + URI uri = + URI.create(String.format("redis://%s:%d", redis.getHost(), redis.getFirstMappedPort())); + pipeline.getOptions().as(TestPipelineOptions.class).setBlockOnRun(false); + + Call.Result throttleResult = + createRequestStream() + .apply( + "throttle", + ThrottleWithExternalResource.usingRedis( + uri, QUOTA_ID, UUID.randomUUID().toString(), QUOTA, REQUEST_CODER)); + + // Assert that we are not getting any errors and successfully emitting throttled elements. + PAssert.that(throttleResult.getFailures()).empty(); + PAssert.thatSingleton( + throttleResult + .getResponses() + .apply( + "window throttled", Window.into(FixedWindows.of(Duration.standardSeconds(1L)))) + .apply( + "count throttled", + Combine.globally(Count.combineFn()).withoutDefaults())) + .notEqualTo(0L); + + // Assert that all the throttled data is not corrupted. + PAssert.that( + throttleResult + .getResponses() + .apply( + "window throttled before extraction", + Window.into(FixedWindows.of(Duration.standardSeconds(1L)))) + .apply( + "extract request id", + MapElements.into(strings()).via(input -> checkStateNotNull(input).getId())) + .apply("distinct", Distinct.create())) + .containsInAnyOrder(QUOTA_ID); + + // Call the Echo service with throttled requests. + Call.Result echoResult = + throttleResult + .getResponses() + .apply("call", Call.ofCallerAndSetupTeardown(client, RESPONSE_CODER)); + + // Assert that there were no errors. + PAssert.that(echoResult.getFailures()).empty(); + + // Assert that the responses match the requests. + PAssert.that( + echoResult + .getResponses() + .apply( + "window responses before extraction", + Window.into(FixedWindows.of(Duration.standardSeconds(1L)))) + .apply( + "extract response id", + MapElements.into(strings()).via(input -> checkStateNotNull(input).getId()))) + .containsInAnyOrder(QUOTA_ID); + + PipelineResult job = pipeline.run(); + job.waitUntilFinish(Duration.standardSeconds(3L)); + } + + private static EchoRequest createRequest() { + return EchoRequest.newBuilder().setId(QUOTA_ID).setPayload(PAYLOAD).build(); + } + + private PCollection createRequestStream() { + return pipeline + .apply("impulse", PeriodicImpulse.create().withInterval(Duration.millis(10L))) + .apply( + "requests", + MapElements.into(TypeDescriptor.of(EchoRequest.class)).via(ignored -> createRequest())); + } +} From ebf23311c114efb1ce3707ae81f913f95f9cd959 Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Mon, 13 Nov 2023 22:02:42 +0000 Subject: [PATCH 6/8] Update javadoc --- .../apache/beam/io/requestresponse/Cache.java | 60 +++++----- .../ThrottleWithExternalResource.java | 109 ++++++++++++++++-- 2 files changed, 132 insertions(+), 37 deletions(-) diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java index 1eb09a9386535..695b1c68bc33e 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java @@ -36,9 +36,10 @@ class Cache { /** - * Read {@link RequestT} {@link ResponseT} associations from a cache. The {@link KV} value is null - * when no association exists. This method does not enforce {@link Coder#verifyDeterministic} and - * defers to the user to determine whether to enforce this given the cache implementation. + * Instantiates a {@link Call} {@link PTransform} that reads {@link RequestT} {@link ResponseT} + * associations from a cache. The {@link KV} value is null when no association exists. This method + * does not enforce {@link Coder#verifyDeterministic} and defers to the user to determine whether + * to enforce this given the cache implementation. */ static < @NonNull RequestT, @@ -57,6 +58,18 @@ class Cache { implementsCallerSetupTeardown, KvCoder.of(requestTCoder, responseTCoder)); } + /** + * Instantiates a {@link Call} {@link PTransform}, calling {@link #read} with a {@link Caller} + * that employs a redis client. + * + *

This method requires both the {@link RequestT} and {@link ResponseT}s' {@link + * Coder#verifyDeterministic}. Otherwise, it throws a {@link NonDeterministicException}. + * + *

Redis is designed for multiple workloads, simultaneously + * reading and writing to a shared instance. See Redis FAQ for more information on important + * considerations when using this method to achieve cache reads. + */ static <@NonNull RequestT, @Nullable ResponseT> PTransform< @NonNull PCollection<@NonNull RequestT>, @@ -94,6 +107,18 @@ class Cache { return Call.ofCallerAndSetupTeardown(implementsCallerSetupTeardown, kvCoder); } + /** + * Instantiates a {@link Call} {@link PTransform}, calling {@link #write} with a {@link Caller} + * that employs a redis client. + * + *

This method requires both the {@link RequestT} and {@link ResponseT}s' {@link + * Coder#verifyDeterministic}. Otherwise, it throws a {@link NonDeterministicException}. + * + *

Redis is designed for multiple workloads, simultaneously + * reading and writing to a shared instance. See Redis FAQ for more information on important + * considerations when using this method to achieve cache writes. + */ static <@NonNull RequestT, @NonNull ResponseT> PTransform< @NonNull PCollection>, @@ -109,25 +134,11 @@ class Cache { KvCoder.of(requestTCoder, responseTCoder)); } - static <@NonNull RequestT, ResponseT> UsingRedis usingRedis( - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder, - @NonNull RedisClient client) - throws Coder.NonDeterministicException { - return new UsingRedis<>(requestTCoder, responseTCoder, client); - } - - /** Provides cache read and write support using a {@link RedisClient}. */ - static class UsingRedis<@NonNull RequestT, ResponseT> { + private static class UsingRedis<@NonNull RequestT, ResponseT> { private final @NonNull Coder<@NonNull RequestT> requestTCoder; private final @NonNull Coder<@Nullable ResponseT> responseTCoder; private final @NonNull RedisClient client; - /** - * Instantiates a {@link UsingRedis}. Given redis reads and writes using bytes, throws a {@link - * Coder.NonDeterministicException} if either {@link RequestT} or {@link ResponseT} {@link - * Coder} fails to {@link Coder#verifyDeterministic}. - */ private UsingRedis( @NonNull Coder<@NonNull RequestT> requestTCoder, @NonNull Coder<@Nullable ResponseT> responseTCoder, @@ -140,21 +151,16 @@ private UsingRedis( this.responseTCoder = responseTCoder; } - /** Instantiates {@link Read}. */ - Read<@NonNull RequestT, @Nullable ResponseT> read() { + private Read<@NonNull RequestT, @Nullable ResponseT> read() { return new Read<>(requestTCoder, responseTCoder, client); } - /** - * Instantiates {@link Write}. The {@link Duration} determines how long the associated {@link - * RequestT} and {@link ResponseT} lasts in the cache. - */ - Write<@NonNull RequestT, @NonNull ResponseT> write(@NonNull Duration expiry) { + private Write<@NonNull RequestT, @NonNull ResponseT> write(@NonNull Duration expiry) { return new Write<>(expiry, requestTCoder, responseTCoder, client); } /** Reads associated {@link RequestT} {@link ResponseT} using a {@link RedisClient}. */ - static class Read<@NonNull RequestT, @Nullable ResponseT> + private static class Read<@NonNull RequestT, @Nullable ResponseT> implements Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>>, SetupTeardown { @@ -203,7 +209,7 @@ public void teardown() throws UserCodeExecutionException { } } - static class Write<@NonNull RequestT, @NonNull ResponseT> + private static class Write<@NonNull RequestT, @NonNull ResponseT> implements Caller< @NonNull KV<@NonNull RequestT, @NonNull ResponseT>, @NonNull KV<@NonNull RequestT, @NonNull ResponseT>>, diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java index 25c94c4c31369..e1440c79a3895 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java @@ -43,21 +43,64 @@ import org.joda.time.Duration; import org.joda.time.Instant; +/** + * Throttles a {@link T} {@link PCollection} using an external resource. + * + *

{@link ThrottleWithExternalResource} makes use of {@link PeriodicImpulse} as it needs to + * coordinate three {@link PTransform}s concurrently. Usage of {@link ThrottleWithExternalResource} + * should consider the impact of {@link PeriodicImpulse} on the pipeline. + * + *

Usage of {@link ThrottleWithExternalResource} is completely optional and serves as one of many + * methods by {@link RequestResponseIO} to protect against API overuse. Usage should not depend on + * {@link ThrottleWithExternalResource} alone to achieve API overuse prevention for several reasons. + * The underlying external resource may not scale at all or as fast as a Beam Runner. The external + * resource itself may be an API with its own quota that {@link ThrottleWithExternalResource} does + * not consider. + * + *

{@link ThrottleWithExternalResource} makes use of several {@link Caller}s that work together + * to achieve its aim of throttling a {@link T} {@link PCollection}. A {@link RefresherT} is a + * {@link Caller} that takes an {@link Instant} and refreshes a shared {@link Quota}. An {@link + * EnqueuerT} enqueues a {@link T} element while a {@link DequeuerT} dequeues said element when the + * {@link ReporterT} reports that the stored {@link Quota#getNumRequests} is >0. Finally, a {@link + * DecrementerT} decrements from the shared {@link Quota} value, additionally reporting the value + * after performing the action. + * + *

{@link ThrottleWithExternalResource} instantiates and applies two {@link Call} {@link + * PTransform}s using the aforementioned {@link Caller}s {@link RefresherT} and {@link EnqueuerT}. + * {@link ThrottleWithExternalResource} calls {@link ReporterT}, {@link DequeuerT}, {@link + * DecrementerT} within its {@link DoFn}, emitting the dequeued {@link T} when the {@link ReporterT} + * reports a value >0. As an additional safety check, the DoFn checks whether the {@link Quota} + * value after {@link DecrementerT}'s action is <0, signaling that multiple workers are attempting + * the same too fast and therefore exists the DoFn allowing for the next refresh. + * + *

{@link ThrottleWithExternalResource} flattens errors emitted from {@link EnqueuerT}, {@link + * RefresherT}, and its own {@link DoFn} into a single {@link ApiIOError} {@link PCollection} that + * is encapsulated, with a {@link T} {@link PCollection} output into a {@link Call.Result}. + */ class ThrottleWithExternalResource< @NonNull T, ReporterT extends Caller<@NonNull String, @NonNull Long> & SetupTeardown, EnqueuerT extends Caller<@NonNull T, Void> & SetupTeardown, DequeuerT extends Caller<@NonNull Instant, @NonNull T> & SetupTeardown, + DecrementerT extends Caller<@NonNull Instant, @NonNull Long> & SetupTeardown, RefresherT extends Caller<@NonNull Instant, Void> & SetupTeardown> extends PTransform<@NonNull PCollection<@NonNull T>, Call.@NonNull Result<@NonNull T>> { - /** Instantiate a {@link ThrottleWithExternalResource} using a {@link RedisClient}. */ + /** + * Instantiate a {@link ThrottleWithExternalResource} using a {@link RedisClient}. + * + *

Redis is designed for multiple workloads, simultaneously + * reading and writing to a shared instance. See Redis FAQ for more information on important + * considerations when using Redis as {@link ThrottleWithExternalResource}'s external resource. + */ static <@NonNull T> ThrottleWithExternalResource< @NonNull T, RedisReporter, RedisEnqueuer<@NonNull T>, - RedisDequeur<@NonNull T>, + RedisDequeuer<@NonNull T>, + RedisDecrementer, RedisRefresher> usingRedis( URI uri, @@ -70,14 +113,16 @@ class ThrottleWithExternalResource< @NonNull T, RedisReporter, RedisEnqueuer<@NonNull T>, - RedisDequeur<@NonNull T>, + RedisDequeuer<@NonNull T>, + RedisDecrementer, RedisRefresher>( quota, quotaIdentifier, coder, new RedisReporter(uri), new RedisEnqueuer<>(uri, queueKey, coder), - new RedisDequeur<>(uri, coder, queueKey), + new RedisDequeuer<>(uri, coder, queueKey), + new RedisDecrementer(uri, queueKey), new RedisRefresher(uri, quota, quotaIdentifier)); } @@ -89,6 +134,7 @@ class ThrottleWithExternalResource< private final @NonNull ReporterT reporterT; private final @NonNull EnqueuerT enqueuerT; private final @NonNull DequeuerT dequeuerT; + private final @NonNull DecrementerT decrementerT; private final @NonNull RefresherT refresherT; ThrottleWithExternalResource( @@ -98,6 +144,7 @@ class ThrottleWithExternalResource< @NonNull ReporterT reporterT, @NonNull EnqueuerT enqueuerT, @NonNull DequeuerT dequeuerT, + @NonNull DecrementerT decrementerT, @NonNull RefresherT refresherT) throws Coder.NonDeterministicException { this.quotaIdentifier = quotaIdentifier; @@ -108,6 +155,7 @@ class ThrottleWithExternalResource< this.coder = coder; this.enqueuerT = enqueuerT; this.dequeuerT = dequeuerT; + this.decrementerT = decrementerT; this.refresherT = refresherT; } @@ -135,7 +183,12 @@ class ThrottleWithExternalResource< "throttle/fn", ParDo.of( new ThrottleFn( - quotaIdentifier, dequeuerT, reporterT, outputTag, failureTag)) + quotaIdentifier, + dequeuerT, + decrementerT, + reporterT, + outputTag, + failureTag)) .withOutputTags(outputTag, TupleTagList.of(failureTag))); PCollection errors = @@ -165,6 +218,7 @@ class ThrottleWithExternalResource< private class ThrottleFn extends DoFn<@NonNull Instant, @NonNull T> { private final String quotaIdentifier; private final DequeuerT dequeuerT; + private final DecrementerT decrementerT; private final ReporterT reporterT; private final TupleTag<@NonNull T> outputTag; private final TupleTag<@NonNull ApiIOError> failureTag; @@ -172,11 +226,13 @@ private class ThrottleFn extends DoFn<@NonNull Instant, @NonNull T> { private ThrottleFn( String quotaIdentifier, DequeuerT dequeuerT, + DecrementerT decrementerT, ReporterT reporterT, TupleTag<@NonNull T> outputTag, TupleTag<@NonNull ApiIOError> failureTag) { this.quotaIdentifier = quotaIdentifier; this.dequeuerT = dequeuerT; + this.decrementerT = decrementerT; this.reporterT = reporterT; this.outputTag = outputTag; this.failureTag = failureTag; @@ -190,16 +246,31 @@ public void process(@Element @NonNull Instant instant, MultiOutputReceiver recei return; } - // Dequeue an element if quota available. + // Decrement the quota. + Long quotaAfterDecrement = decrementerT.call(instant); + + // As an additional protection we check what the quota is after decrementing. A value + // < 0 signals that multiple simultaneous workers have attempted to decrement too quickly. + // We don't bother adding the quota back to prevent additional workers from doing the same + // and just wait for the next refresh, exiting the DoFn. + if (quotaAfterDecrement < 0) { + return; + } + + // Dequeue an element if quota available. An error here would not result in loss of data + // as no element would successfully dequeue from the external resource but instead throw. T element = dequeuerT.call(instant); - // Emit element. + + // Finally, emit the element. receiver.get(outputTag).output(element); + } catch (UserCodeExecutionException e) { receiver .get(failureTag) .output( ApiIOError.builder() - .setRequestAsJsonString(String.format("{\"request\": \"%s\"}", instant)) + // no request to emit as part of the error. + .setRequestAsJsonString("") .setMessage(Optional.ofNullable(e.getMessage()).orElse("")) .setObservedTimestamp(Instant.now()) .setStackTrace(Throwables.getStackTraceAsString(e)) @@ -211,6 +282,7 @@ public void process(@Element @NonNull Instant instant, MultiOutputReceiver recei public void setup() throws UserCodeExecutionException { enqueuerT.setup(); dequeuerT.setup(); + decrementerT.setup(); reporterT.setup(); } @@ -218,6 +290,7 @@ public void setup() throws UserCodeExecutionException { public void teardown() throws UserCodeExecutionException { enqueuerT.teardown(); dequeuerT.teardown(); + decrementerT.teardown(); reporterT.teardown(); } } @@ -258,13 +331,13 @@ public Void call(@NonNull T request) throws UserCodeExecutionException { } } - private static class RedisDequeur<@NonNull T> extends RedisSetupTeardown + private static class RedisDequeuer<@NonNull T> extends RedisSetupTeardown implements Caller<@NonNull Instant, @NonNull T> { private final Coder<@NonNull T> coder; private final String key; - private RedisDequeur(URI uri, Coder<@NonNull T> coder, String key) { + private RedisDequeuer(URI uri, Coder<@NonNull T> coder, String key) { super(new RedisClient(uri)); this.coder = coder; this.key = key; @@ -282,6 +355,22 @@ private RedisDequeur(URI uri, Coder<@NonNull T> coder, String key) { } } + private static class RedisDecrementer extends RedisSetupTeardown + implements Caller<@NonNull Instant, @NonNull Long> { + + private final String key; + + private RedisDecrementer(URI uri, String key) { + super(new RedisClient(uri)); + this.key = key; + } + + @Override + public @NonNull Long call(@NonNull Instant request) throws UserCodeExecutionException { + return client.decr(key); + } + } + private static class RedisRefresher extends RedisSetupTeardown implements Caller<@NonNull Instant, Void> { private final Quota quota; From 9620eeedea20a9f7016fb5be7f9b57181e5cd01f Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Sun, 19 Nov 2023 21:32:09 -0800 Subject: [PATCH 7/8] Edit per PR comments --- .../apache/beam/io/requestresponse/Cache.java | 126 +++++++--------- .../ThrottleWithExternalResource.java | 139 ++++++++---------- .../{CacheTestIT.java => CacheIT.java} | 2 +- .../beam/io/requestresponse/CallTest.java | 8 +- ...=> EchoGRPCCallerWithSetupTeardownIT.java} | 6 +- ...allerTestIT.java => EchoHTTPCallerIT.java} | 8 +- .../io/requestresponse/EchoITOptions.java | 6 +- ...isClientTestIT.java => RedisClientIT.java} | 2 +- ...va => ThrottleWithExternalResourceIT.java} | 8 +- 9 files changed, 133 insertions(+), 172 deletions(-) rename sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/{CacheTestIT.java => CacheIT.java} (99%) rename sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/{EchoGRPCCallerWithSetupTeardownTestIT.java => EchoGRPCCallerWithSetupTeardownIT.java} (97%) rename sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/{EchoHTTPCallerTestIT.java => EchoHTTPCallerIT.java} (94%) rename sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/{RedisClientTestIT.java => RedisClientIT.java} (99%) rename sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/{ThrottleWithExternalResourceTestIT.java => ThrottleWithExternalResourceIT.java} (96%) diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java index 695b1c68bc33e..b8e526c4829b0 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Cache.java @@ -28,12 +28,11 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; -import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; /** Transforms for reading and writing request/response associations to a cache. */ -class Cache { +final class Cache { /** * Instantiates a {@link Call} {@link PTransform} that reads {@link RequestT} {@link ResponseT} @@ -42,18 +41,14 @@ class Cache { * to enforce this given the cache implementation. */ static < - @NonNull RequestT, + RequestT, @Nullable ResponseT, CallerSetupTeardownT extends - @NonNull Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>> - & SetupTeardown> - PTransform< - @NonNull PCollection<@NonNull RequestT>, - Call.@NonNull Result>> - read( - @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder) { + Caller> & SetupTeardown> + PTransform, Call.Result>> read( + CallerSetupTeardownT implementsCallerSetupTeardown, + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder) { return Call.ofCallerAndSetupTeardown( implementsCallerSetupTeardown, KvCoder.of(requestTCoder, responseTCoder)); } @@ -70,14 +65,12 @@ class Cache { * href="https://redis.io/docs/get-started/faq/">Redis FAQ for more information on important * considerations when using this method to achieve cache reads. */ - static <@NonNull RequestT, @Nullable ResponseT> - PTransform< - @NonNull PCollection<@NonNull RequestT>, - Call.@NonNull Result>> + static + PTransform, Call.Result>> readUsingRedis( - @NonNull RedisClient client, - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder) + RedisClient client, + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder) throws NonDeterministicException { return read( new UsingRedis<>(requestTCoder, responseTCoder, client).read(), @@ -91,19 +84,13 @@ class Cache { * given the cache implementation. */ static < - @NonNull RequestT, - @NonNull ResponseT, + RequestT, + ResponseT, CallerSetupTeardownT extends - Caller< - KV<@NonNull RequestT, @NonNull ResponseT>, - KV<@NonNull RequestT, @NonNull ResponseT>> - & SetupTeardown> - PTransform< - @NonNull PCollection>, - Call.@NonNull Result>> - write( - @NonNull CallerSetupTeardownT implementsCallerSetupTeardown, - @NonNull KvCoder<@NonNull RequestT, @NonNull ResponseT> kvCoder) { + Caller, KV> & SetupTeardown> + PTransform>, Call.Result>> write( + CallerSetupTeardownT implementsCallerSetupTeardown, + KvCoder kvCoder) { return Call.ofCallerAndSetupTeardown(implementsCallerSetupTeardown, kvCoder); } @@ -119,30 +106,28 @@ class Cache { * href="https://redis.io/docs/get-started/faq/">Redis FAQ for more information on important * considerations when using this method to achieve cache writes. */ - static <@NonNull RequestT, @NonNull ResponseT> - PTransform< - @NonNull PCollection>, - Call.@NonNull Result>> + static + PTransform>, Call.Result>> writeUsingRedis( - @NonNull Duration expiry, - @NonNull RedisClient client, - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder) + Duration expiry, + RedisClient client, + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder) throws NonDeterministicException { return write( new UsingRedis<>(requestTCoder, responseTCoder, client).write(expiry), KvCoder.of(requestTCoder, responseTCoder)); } - private static class UsingRedis<@NonNull RequestT, ResponseT> { - private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@Nullable ResponseT> responseTCoder; - private final @NonNull RedisClient client; + private static class UsingRedis { + private final Coder requestTCoder; + private final Coder<@Nullable ResponseT> responseTCoder; + private final RedisClient client; private UsingRedis( - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder, - @NonNull RedisClient client) + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder, + RedisClient client) throws Coder.NonDeterministicException { this.client = client; requestTCoder.verifyDeterministic(); @@ -151,34 +136,33 @@ private UsingRedis( this.responseTCoder = responseTCoder; } - private Read<@NonNull RequestT, @Nullable ResponseT> read() { + private Read read() { return new Read<>(requestTCoder, responseTCoder, client); } - private Write<@NonNull RequestT, @NonNull ResponseT> write(@NonNull Duration expiry) { + private Write write(Duration expiry) { return new Write<>(expiry, requestTCoder, responseTCoder, client); } /** Reads associated {@link RequestT} {@link ResponseT} using a {@link RedisClient}. */ - private static class Read<@NonNull RequestT, @Nullable ResponseT> - implements Caller<@NonNull RequestT, KV<@NonNull RequestT, @Nullable ResponseT>>, - SetupTeardown { + private static class Read + implements Caller>, SetupTeardown { - private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@Nullable ResponseT> responseTCoder; - private final @NonNull RedisClient client; + private final Coder requestTCoder; + private final Coder<@Nullable ResponseT> responseTCoder; + private final RedisClient client; private Read( - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder, - @NonNull RedisClient client) { + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder, + RedisClient client) { this.requestTCoder = requestTCoder; this.responseTCoder = responseTCoder; this.client = client; } @Override - public KV<@NonNull RequestT, @Nullable ResponseT> call(@NonNull RequestT request) + public KV call(RequestT request) throws UserCodeExecutionException { ByteArrayOutputStream baos = new ByteArrayOutputStream(); try { @@ -209,21 +193,18 @@ public void teardown() throws UserCodeExecutionException { } } - private static class Write<@NonNull RequestT, @NonNull ResponseT> - implements Caller< - @NonNull KV<@NonNull RequestT, @NonNull ResponseT>, - @NonNull KV<@NonNull RequestT, @NonNull ResponseT>>, - SetupTeardown { - private final @NonNull Duration expiry; - private final @NonNull Coder<@NonNull RequestT> requestTCoder; - private final @NonNull Coder<@Nullable ResponseT> responseTCoder; - private final @NonNull RedisClient client; + private static class Write + implements Caller, KV>, SetupTeardown { + private final Duration expiry; + private final Coder requestTCoder; + private final Coder<@Nullable ResponseT> responseTCoder; + private final RedisClient client; private Write( - @NonNull Duration expiry, - @NonNull Coder<@NonNull RequestT> requestTCoder, - @NonNull Coder<@Nullable ResponseT> responseTCoder, - @NonNull RedisClient client) { + Duration expiry, + Coder requestTCoder, + Coder<@Nullable ResponseT> responseTCoder, + RedisClient client) { this.expiry = expiry; this.requestTCoder = requestTCoder; this.responseTCoder = responseTCoder; @@ -231,8 +212,7 @@ private Write( } @Override - public @NonNull KV<@NonNull RequestT, @NonNull ResponseT> call( - @NonNull KV<@NonNull RequestT, @NonNull ResponseT> request) + public KV call(KV request) throws UserCodeExecutionException { ByteArrayOutputStream keyStream = new ByteArrayOutputStream(); ByteArrayOutputStream valueStream = new ByteArrayOutputStream(); diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java index e1440c79a3895..7b5c471a5ad21 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java @@ -39,7 +39,6 @@ import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteSource; -import org.checkerframework.checker.nullness.qual.NonNull; import org.joda.time.Duration; import org.joda.time.Instant; @@ -78,13 +77,13 @@ * is encapsulated, with a {@link T} {@link PCollection} output into a {@link Call.Result}. */ class ThrottleWithExternalResource< - @NonNull T, - ReporterT extends Caller<@NonNull String, @NonNull Long> & SetupTeardown, - EnqueuerT extends Caller<@NonNull T, Void> & SetupTeardown, - DequeuerT extends Caller<@NonNull Instant, @NonNull T> & SetupTeardown, - DecrementerT extends Caller<@NonNull Instant, @NonNull Long> & SetupTeardown, - RefresherT extends Caller<@NonNull Instant, Void> & SetupTeardown> - extends PTransform<@NonNull PCollection<@NonNull T>, Call.@NonNull Result<@NonNull T>> { + T, + ReporterT extends Caller & SetupTeardown, + EnqueuerT extends Caller & SetupTeardown, + DequeuerT extends Caller & SetupTeardown, + DecrementerT extends Caller & SetupTeardown, + RefresherT extends Caller & SetupTeardown> + extends PTransform, Call.Result> { /** * Instantiate a {@link ThrottleWithExternalResource} using a {@link RedisClient}. @@ -94,28 +93,18 @@ class ThrottleWithExternalResource< * href="https://redis.io/docs/get-started/faq/">Redis FAQ for more information on important * considerations when using Redis as {@link ThrottleWithExternalResource}'s external resource. */ - static <@NonNull T> + static ThrottleWithExternalResource< - @NonNull T, + T, RedisReporter, - RedisEnqueuer<@NonNull T>, - RedisDequeuer<@NonNull T>, + RedisEnqueuer, + RedisDequeuer, RedisDecrementer, RedisRefresher> - usingRedis( - URI uri, - String quotaIdentifier, - String queueKey, - Quota quota, - Coder<@NonNull T> coder) + usingRedis(URI uri, String quotaIdentifier, String queueKey, Quota quota, Coder coder) throws Coder.NonDeterministicException { return new ThrottleWithExternalResource< - @NonNull T, - RedisReporter, - RedisEnqueuer<@NonNull T>, - RedisDequeuer<@NonNull T>, - RedisDecrementer, - RedisRefresher>( + T, RedisReporter, RedisEnqueuer, RedisDequeuer, RedisDecrementer, RedisRefresher>( quota, quotaIdentifier, coder, @@ -128,24 +117,24 @@ class ThrottleWithExternalResource< private static final Duration THROTTLE_INTERVAL = Duration.standardSeconds(1L); - private final @NonNull Quota quota; - private final @NonNull String quotaIdentifier; - private final @NonNull Coder<@NonNull T> coder; - private final @NonNull ReporterT reporterT; - private final @NonNull EnqueuerT enqueuerT; - private final @NonNull DequeuerT dequeuerT; - private final @NonNull DecrementerT decrementerT; - private final @NonNull RefresherT refresherT; + private final Quota quota; + private final String quotaIdentifier; + private final Coder coder; + private final ReporterT reporterT; + private final EnqueuerT enqueuerT; + private final DequeuerT dequeuerT; + private final DecrementerT decrementerT; + private final RefresherT refresherT; ThrottleWithExternalResource( - @NonNull Quota quota, - @NonNull String quotaIdentifier, - @NonNull Coder<@NonNull T> coder, - @NonNull ReporterT reporterT, - @NonNull EnqueuerT enqueuerT, - @NonNull DequeuerT dequeuerT, - @NonNull DecrementerT decrementerT, - @NonNull RefresherT refresherT) + Quota quota, + String quotaIdentifier, + Coder coder, + ReporterT reporterT, + EnqueuerT enqueuerT, + DequeuerT dequeuerT, + DecrementerT decrementerT, + RefresherT refresherT) throws Coder.NonDeterministicException { this.quotaIdentifier = quotaIdentifier; this.reporterT = reporterT; @@ -160,27 +149,27 @@ class ThrottleWithExternalResource< } @Override - public Call.@NonNull Result<@NonNull T> expand(PCollection<@NonNull T> input) { + public Call.Result expand(PCollection input) { Pipeline pipeline = input.getPipeline(); // Refresh known quota to control the throttle rate. Call.Result refreshResult = pipeline - .apply("quota/impulse", PeriodicImpulse.create().withInterval(quota.getInterval())) - .apply("quota/refresh", getRefresher()); + .apply("quota impulse", PeriodicImpulse.create().withInterval(quota.getInterval())) + .apply("quota refresh", getRefresher()); // Enqueue T elements. Call.Result enqueuResult = input.apply("enqueue", getEnqueuer()); - TupleTag<@NonNull T> outputTag = new TupleTag<@NonNull T>() {}; - TupleTag<@NonNull ApiIOError> failureTag = new TupleTag<@NonNull ApiIOError>() {}; + TupleTag outputTag = new TupleTag() {}; + TupleTag failureTag = new TupleTag() {}; // Perform Throttle. PCollectionTuple pct = pipeline - .apply("throttle/impulse", PeriodicImpulse.create().withInterval(THROTTLE_INTERVAL)) + .apply("throttle impulse", PeriodicImpulse.create().withInterval(THROTTLE_INTERVAL)) .apply( - "throttle/fn", + "throttle fn", ParDo.of( new ThrottleFn( quotaIdentifier, @@ -195,41 +184,41 @@ class ThrottleWithExternalResource< PCollectionList.of(refreshResult.getFailures()) .and(enqueuResult.getFailures()) .and(pct.get(failureTag)) - .apply("errors/flatten", Flatten.pCollections()); + .apply("errors flatten", Flatten.pCollections()); - TupleTag<@NonNull T> resultOutputTag = new TupleTag<@NonNull T>() {}; - TupleTag<@NonNull ApiIOError> resultFailureTag = new TupleTag<@NonNull ApiIOError>() {}; + TupleTag resultOutputTag = new TupleTag() {}; + TupleTag resultFailureTag = new TupleTag() {}; - return Call.Result.<@NonNull T>of( + return Call.Result.of( coder, resultOutputTag, resultFailureTag, PCollectionTuple.of(resultOutputTag, pct.get(outputTag)).and(resultFailureTag, errors)); } - private Call<@NonNull Instant, Void> getRefresher() { + private Call getRefresher() { return Call.ofCallerAndSetupTeardown(refresherT, VoidCoder.of()); } - private Call<@NonNull T, Void> getEnqueuer() { + private Call getEnqueuer() { return Call.ofCallerAndSetupTeardown(enqueuerT, VoidCoder.of()); } - private class ThrottleFn extends DoFn<@NonNull Instant, @NonNull T> { + private class ThrottleFn extends DoFn { private final String quotaIdentifier; private final DequeuerT dequeuerT; private final DecrementerT decrementerT; private final ReporterT reporterT; - private final TupleTag<@NonNull T> outputTag; - private final TupleTag<@NonNull ApiIOError> failureTag; + private final TupleTag outputTag; + private final TupleTag failureTag; private ThrottleFn( String quotaIdentifier, DequeuerT dequeuerT, DecrementerT decrementerT, ReporterT reporterT, - TupleTag<@NonNull T> outputTag, - TupleTag<@NonNull ApiIOError> failureTag) { + TupleTag outputTag, + TupleTag failureTag) { this.quotaIdentifier = quotaIdentifier; this.dequeuerT = dequeuerT; this.decrementerT = decrementerT; @@ -239,7 +228,7 @@ private ThrottleFn( } @ProcessElement - public void process(@Element @NonNull Instant instant, MultiOutputReceiver receiver) { + public void process(@Element Instant instant, MultiOutputReceiver receiver) { // Check for available quota. try { if (reporterT.call(quotaIdentifier) <= 0L) { @@ -295,31 +284,29 @@ public void teardown() throws UserCodeExecutionException { } } - private static class RedisReporter extends RedisSetupTeardown - implements Caller<@NonNull String, @NonNull Long> { + private static class RedisReporter extends RedisSetupTeardown implements Caller { private RedisReporter(URI uri) { super(new RedisClient(uri)); } @Override - public @NonNull Long call(@NonNull String request) throws UserCodeExecutionException { + public Long call(String request) throws UserCodeExecutionException { return client.getLong(request); } } - private static class RedisEnqueuer<@NonNull T> extends RedisSetupTeardown - implements Caller<@NonNull T, Void> { + private static class RedisEnqueuer extends RedisSetupTeardown implements Caller { private final String key; - private final Coder<@NonNull T> coder; + private final Coder coder; - private RedisEnqueuer(URI uri, String key, Coder<@NonNull T> coder) { + private RedisEnqueuer(URI uri, String key, Coder coder) { super(new RedisClient(uri)); this.key = key; this.coder = coder; } @Override - public Void call(@NonNull T request) throws UserCodeExecutionException { + public Void call(T request) throws UserCodeExecutionException { ByteArrayOutputStream baos = new ByteArrayOutputStream(); try { coder.encode(request, baos); @@ -331,20 +318,19 @@ public Void call(@NonNull T request) throws UserCodeExecutionException { } } - private static class RedisDequeuer<@NonNull T> extends RedisSetupTeardown - implements Caller<@NonNull Instant, @NonNull T> { + private static class RedisDequeuer extends RedisSetupTeardown implements Caller { - private final Coder<@NonNull T> coder; + private final Coder coder; private final String key; - private RedisDequeuer(URI uri, Coder<@NonNull T> coder, String key) { + private RedisDequeuer(URI uri, Coder coder, String key) { super(new RedisClient(uri)); this.coder = coder; this.key = key; } @Override - public @NonNull T call(@NonNull Instant request) throws UserCodeExecutionException { + public T call(Instant request) throws UserCodeExecutionException { byte[] bytes = client.lpop(key); try { return checkStateNotNull(coder.decode(ByteSource.wrap(bytes).openStream())); @@ -356,7 +342,7 @@ private RedisDequeuer(URI uri, Coder<@NonNull T> coder, String key) { } private static class RedisDecrementer extends RedisSetupTeardown - implements Caller<@NonNull Instant, @NonNull Long> { + implements Caller { private final String key; @@ -366,13 +352,12 @@ private RedisDecrementer(URI uri, String key) { } @Override - public @NonNull Long call(@NonNull Instant request) throws UserCodeExecutionException { + public Long call(Instant request) throws UserCodeExecutionException { return client.decr(key); } } - private static class RedisRefresher extends RedisSetupTeardown - implements Caller<@NonNull Instant, Void> { + private static class RedisRefresher extends RedisSetupTeardown implements Caller { private final Quota quota; private final String key; @@ -383,7 +368,7 @@ private RedisRefresher(URI uri, Quota quota, String key) { } @Override - public Void call(@NonNull Instant request) throws UserCodeExecutionException { + public Void call(Instant request) throws UserCodeExecutionException { client.setex(key, quota.getNumRequests(), quota.getInterval()); return null; } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheIT.java similarity index 99% rename from sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java rename to sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheIT.java index fed09a2a19041..95497e6013af1 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CacheIT.java @@ -39,7 +39,7 @@ /** Integration tests for {@link Cache}. */ @RunWith(JUnit4.class) -public class CacheTestIT { +public class CacheIT { @Rule public TestPipeline writePipeline = TestPipeline.create(); @Rule public TestPipeline readPipeline = TestPipeline.create(); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java index 1314c1b772903..1566d1725295e 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java @@ -43,10 +43,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; -import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; -import org.checkerframework.checker.nullness.qual.UnknownKeyFor; import org.jetbrains.annotations.NotNull; import org.joda.time.Duration; import org.junit.Rule; @@ -535,8 +533,7 @@ public void encode(Request value, @NotNull OutputStream outStream) } @Override - public void verifyDeterministic() - throws @UnknownKeyFor @NonNull @Initialized NonDeterministicException { + public void verifyDeterministic() throws NonDeterministicException { ID_CODER.verifyDeterministic(); } } @@ -565,8 +562,7 @@ public Response decode(@NotNull InputStream inStream) throws CoderException, IOE } @Override - public void verifyDeterministic() - throws @UnknownKeyFor @NonNull @Initialized NonDeterministicException { + public void verifyDeterministic() throws NonDeterministicException { ID_CODER.verifyDeterministic(); } } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java similarity index 97% rename from sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownTestIT.java rename to sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java index 14b6e9e6433d4..ea25c76433376 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java @@ -41,7 +41,7 @@ * running integration tests. */ @RunWith(JUnit4.class) -public class EchoGRPCCallerWithSetupTeardownTestIT { +public class EchoGRPCCallerWithSetupTeardownIT { private static @MonotonicNonNull EchoITOptions options; private static @MonotonicNonNull EchoGRPCCallerWithSetupTeardown client; @@ -50,11 +50,11 @@ public class EchoGRPCCallerWithSetupTeardownTestIT { @BeforeClass public static void setUp() throws UserCodeExecutionException { options = readIOTestPipelineOptions(EchoITOptions.class); - if (options.getgRPCEndpointAddress().isEmpty()) { + if (options.getGrpcEndpointAddress().isEmpty()) { throw new RuntimeException( "--gRPCEndpointAddress is missing. See " + EchoITOptions.class + "for details."); } - client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getgRPCEndpointAddress())); + client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getGrpcEndpointAddress())); checkStateNotNull(client).setup(); EchoRequest request = createShouldExceedQuotaRequest(); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java similarity index 94% rename from sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerTestIT.java rename to sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java index fa0cb93781100..5f13a1eb094e0 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java @@ -36,12 +36,12 @@ import org.junit.runners.JUnit4; /** - * Tests for {@link EchoHTTPCallerTestIT} on a deployed {@link EchoServiceGrpc} instance's HTTP - * handler. See {@link EchoITOptions} for details on the required parameters and how to provide - * these for running integration tests. + * Tests for {@link EchoHTTPCallerIT} on a deployed {@link EchoServiceGrpc} instance's HTTP handler. + * See {@link EchoITOptions} for details on the required parameters and how to provide these for + * running integration tests. */ @RunWith(JUnit4.class) -public class EchoHTTPCallerTestIT { +public class EchoHTTPCallerIT { private static @MonotonicNonNull EchoITOptions options; private static @MonotonicNonNull EchoHTTPCaller client; diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java index 0219d5d4f6341..bd46c52e1a6cd 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java @@ -31,16 +31,16 @@ * * 

  *   ./gradlew :sdks:java:io:rrio:integrationTest -DintegrationTestPipelineOptions='[
- *      "--gRPCEndpointAddress=",
+ *      "--grpcEndpointAddress=",
  *      "--httpEndpointAddress="
  *   ]'
  *
*/ public interface EchoITOptions extends PipelineOptions { @Description("The gRPC address of the Echo API endpoint, typically of the form :.") - String getgRPCEndpointAddress(); + String getGrpcEndpointAddress(); - void setgRPCEndpointAddress(String value); + void setGrpcEndpointAddress(String value); @Description("The HTTP address of the Echo API endpoint; must being with http(s)://") String getHttpEndpointAddress(); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientIT.java similarity index 99% rename from sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java rename to sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientIT.java index 16b83b1637afa..939515836bff3 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/RedisClientIT.java @@ -50,7 +50,7 @@ /** Integration tests for {@link RedisClient}. */ @RunWith(JUnit4.class) -public class RedisClientTestIT { +public class RedisClientIT { private static final String CONTAINER_IMAGE_NAME = "redis:5.0.3-alpine"; private static final Integer PORT = 6379; diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java similarity index 96% rename from sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java rename to sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java index 007e418f21393..afd590e28287b 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceTestIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java @@ -55,7 +55,7 @@ * Integration tests for {@link ThrottleWithExternalResource}. See {@link EchoITOptions} for details * on the required parameters and how to provide these for running integration tests. */ -public class ThrottleWithExternalResourceTestIT { +public class ThrottleWithExternalResourceIT { @Rule public TestPipeline pipeline = TestPipeline.create(); @@ -77,11 +77,11 @@ public class ThrottleWithExternalResourceTestIT { @BeforeClass public static void setUp() throws UserCodeExecutionException { options = readIOTestPipelineOptions(EchoITOptions.class); - if (options.getgRPCEndpointAddress().isEmpty()) { + if (options.getGrpcEndpointAddress() == null || options.getGrpcEndpointAddress().isEmpty()) { throw new RuntimeException( - "--gRPCEndpointAddress is missing. See " + EchoITOptions.class + "for details."); + "--grpcEndpointAddress is missing. See " + EchoITOptions.class + "for details."); } - client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getgRPCEndpointAddress())); + client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getGrpcEndpointAddress())); checkStateNotNull(client).setup(); try { From f63c13f0c96fb4b62f331bbcdc4db67916e5a08f Mon Sep 17 00:00:00 2001 From: Damon Douglas Date: Wed, 22 Nov 2023 18:34:11 +0000 Subject: [PATCH 8/8] Refacter per PR comments --- .../apache/beam/io/requestresponse/Call.java | 5 ++- .../ThrottleWithExternalResource.java | 32 ++++++++++++++++--- .../EchoGRPCCallerWithSetupTeardownIT.java | 10 ++++-- .../io/requestresponse/EchoHTTPCallerIT.java | 10 ++++-- .../io/requestresponse/EchoITOptions.java | 3 ++ .../ThrottleWithExternalResourceIT.java | 10 ++++-- 6 files changed, 59 insertions(+), 11 deletions(-) diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java index d6d0881c11f9d..d52ca971ca47a 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java @@ -100,6 +100,10 @@ Call ofCallerAndSetupTeardown( .build()); } + // TupleTags need to be instantiated for each Call instance. We cannot use a shared + // static instance that is shared for multiple PCollectionTuples when Call is + // instantiated multiple times as it is reused throughout code in this library. + private final TupleTag responseTag = new TupleTag() {}; private final TupleTag failureTag = new TupleTag() {}; private final Configuration configuration; @@ -128,7 +132,6 @@ Call withTimeout(Duration timeout) { @Override public @NonNull Result expand(PCollection input) { - TupleTag responseTag = new TupleTag() {}; PCollectionTuple pct = input.apply( diff --git a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java index 7b5c471a5ad21..dffc034770aa7 100644 --- a/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java +++ b/sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResource.java @@ -23,6 +23,8 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.net.URI; +import java.util.ArrayList; +import java.util.List; import java.util.Optional; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.Coder; @@ -277,10 +279,32 @@ public void setup() throws UserCodeExecutionException { @Teardown public void teardown() throws UserCodeExecutionException { - enqueuerT.teardown(); - dequeuerT.teardown(); - decrementerT.teardown(); - reporterT.teardown(); + List messages = new ArrayList<>(); + String format = "%s encountered error during teardown: %s"; + try { + enqueuerT.teardown(); + } catch (UserCodeExecutionException e) { + messages.add(String.format(format, "enqueuerT", e)); + } + try { + dequeuerT.teardown(); + } catch (UserCodeExecutionException e) { + messages.add(String.format(format, "dequeuerT", e)); + } + try { + decrementerT.teardown(); + } catch (UserCodeExecutionException e) { + messages.add(String.format(format, "decrementerT", e)); + } + try { + reporterT.teardown(); + } catch (UserCodeExecutionException e) { + messages.add(String.format(format, "reporterT", e)); + } + + if (!messages.isEmpty()) { + throw new UserCodeExecutionException(String.join("; ", messages)); + } } } diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java index ea25c76433376..c10b7ee1609e5 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoGRPCCallerWithSetupTeardownIT.java @@ -17,6 +17,7 @@ */ package org.apache.beam.io.requestresponse; +import static org.apache.beam.io.requestresponse.EchoITOptions.GRPC_ENDPOINT_ADDRESS_NAME; import static org.apache.beam.sdk.io.common.IOITHelper.readIOTestPipelineOptions; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.junit.Assert.assertEquals; @@ -27,6 +28,7 @@ import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoRequest; import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoResponse; import org.apache.beam.testinfra.mockapis.echo.v1.EchoServiceGrpc; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.NonNull; import org.junit.AfterClass; @@ -50,9 +52,13 @@ public class EchoGRPCCallerWithSetupTeardownIT { @BeforeClass public static void setUp() throws UserCodeExecutionException { options = readIOTestPipelineOptions(EchoITOptions.class); - if (options.getGrpcEndpointAddress().isEmpty()) { + if (Strings.isNullOrEmpty(options.getGrpcEndpointAddress())) { throw new RuntimeException( - "--gRPCEndpointAddress is missing. See " + EchoITOptions.class + "for details."); + "--" + + GRPC_ENDPOINT_ADDRESS_NAME + + " is missing. See " + + EchoITOptions.class + + "for details."); } client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getGrpcEndpointAddress())); checkStateNotNull(client).setup(); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java index 5f13a1eb094e0..10b92b2610d92 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoHTTPCallerIT.java @@ -17,6 +17,7 @@ */ package org.apache.beam.io.requestresponse; +import static org.apache.beam.io.requestresponse.EchoITOptions.HTTP_ENDPOINT_ADDRESS_NAME; import static org.apache.beam.sdk.io.common.IOITHelper.readIOTestPipelineOptions; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.junit.Assert.assertEquals; @@ -28,6 +29,7 @@ import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoRequest; import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoResponse; import org.apache.beam.testinfra.mockapis.echo.v1.EchoServiceGrpc; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.NonNull; import org.junit.BeforeClass; @@ -50,9 +52,13 @@ public class EchoHTTPCallerIT { @BeforeClass public static void setUp() throws UserCodeExecutionException { options = readIOTestPipelineOptions(EchoITOptions.class); - if (options.getHttpEndpointAddress().isEmpty()) { + if (Strings.isNullOrEmpty(options.getHttpEndpointAddress())) { throw new RuntimeException( - "--httpEndpointAddress is missing. See " + EchoITOptions.class + "for details."); + "--" + + HTTP_ENDPOINT_ADDRESS_NAME + + " is missing. See " + + EchoITOptions.class + + "for details."); } client = EchoHTTPCaller.of(URI.create(options.getHttpEndpointAddress())); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java index bd46c52e1a6cd..dabec75089263 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/EchoITOptions.java @@ -37,6 +37,9 @@ * */ public interface EchoITOptions extends PipelineOptions { + String GRPC_ENDPOINT_ADDRESS_NAME = "grpcEndpointAddress"; + String HTTP_ENDPOINT_ADDRESS_NAME = "httpEndpointAddress"; + @Description("The gRPC address of the Echo API endpoint, typically of the form :.") String getGrpcEndpointAddress(); diff --git a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java index afd590e28287b..24db38f926eed 100644 --- a/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java +++ b/sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/ThrottleWithExternalResourceIT.java @@ -17,6 +17,7 @@ */ package org.apache.beam.io.requestresponse; +import static org.apache.beam.io.requestresponse.EchoITOptions.GRPC_ENDPOINT_ADDRESS_NAME; import static org.apache.beam.sdk.io.common.IOITHelper.readIOTestPipelineOptions; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.sdk.values.TypeDescriptors.strings; @@ -42,6 +43,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoRequest; import org.apache.beam.testinfra.mockapis.echo.v1.Echo.EchoResponse; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.joda.time.Duration; import org.junit.AfterClass; @@ -77,9 +79,13 @@ public class ThrottleWithExternalResourceIT { @BeforeClass public static void setUp() throws UserCodeExecutionException { options = readIOTestPipelineOptions(EchoITOptions.class); - if (options.getGrpcEndpointAddress() == null || options.getGrpcEndpointAddress().isEmpty()) { + if (Strings.isNullOrEmpty(options.getGrpcEndpointAddress())) { throw new RuntimeException( - "--grpcEndpointAddress is missing. See " + EchoITOptions.class + "for details."); + "--" + + GRPC_ENDPOINT_ADDRESS_NAME + + " is missing. See " + + EchoITOptions.class + + "for details."); } client = EchoGRPCCallerWithSetupTeardown.of(URI.create(options.getGrpcEndpointAddress())); checkStateNotNull(client).setup();