gguptp commented on code in PR #261: URL: https://github.com/apache/flink-connector-aws/pull/261#discussion_r4505115262
########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/util/DynamoDbStreamsAsyncClientUtil.java: ########## @@ -0,0 +1,104 @@ +/* + * 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.flink.connector.dynamodb.source.util; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.aws.config.AWSConfigConstants; +import org.apache.flink.connector.aws.util.AWSClientUtil; +import org.apache.flink.connector.aws.util.AWSGeneralUtil; +import org.apache.flink.connector.dynamodb.source.config.DynamodbStreamsSourceConfigConstants; + +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption; +import software.amazon.awssdk.http.Protocol; +import software.amazon.awssdk.http.async.SdkAsyncHttpClient; +import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient; +import software.amazon.awssdk.retries.api.RetryStrategy; +import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsAsyncClient; +import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsAsyncClientBuilder; +import software.amazon.awssdk.utils.AttributeMap; + +import java.net.URI; +import java.util.Properties; + +/** + * Creates the {@link DynamoDbStreamsAsyncClient} that the source readers use for non-blocking + * GetRecords calls. The client has the same settings as the {@code DynamoDbStreamsClient} that + * {@link AWSClientUtil#createAwsSyncClient} creates for blocking calls. + */ +@Internal +public final class DynamoDbStreamsAsyncClientUtil { + + /** Same suffix that {@link AWSClientUtil} adds to the default user agent prefix. */ + private static final String V2_USER_AGENT_SUFFIX = " V2"; + + private DynamoDbStreamsAsyncClientUtil() {} + + /** + * Creates a Netty HTTP client with the same settings as the HTTP client for blocking calls. + * + * <p>{@link AWSGeneralUtil#createAsyncHttpClient} uses HTTP/2 by default. DynamoDB Streams + * endpoints use only HTTP/1.1, and an HTTP/2 client fails with "First received frame was not + * SETTINGS". A protocol set on the builder has precedence over the defaults of the helper, so + * the client always uses HTTP/1.1. + */ + public static SdkAsyncHttpClient createAsyncHttpClient() { + return AWSGeneralUtil.createAsyncHttpClient( + AttributeMap.empty(), NettyNioAsyncHttpClient.builder().protocol(Protocol.HTTP1_1)); + } + + /** + * Creates the client. The client gets the region, credentials, endpoint override, and user + * agent prefix from {@code clientProperties}, as {@link AWSClientUtil#createAwsSyncClient} + * does. This method does not use {@link AWSClientUtil#createAwsAsyncClient}, because that + * method replaces the client override configuration and drops the retry strategy. + */ + public static DynamoDbStreamsAsyncClient createAsyncClient( Review Comment: done in this PR. added a `createAwsAsyncClient` overload to `AWSClientUtil` that takes a `ClientOverrideConfiguration.Builder`, like the sync one, with a test in `AWSClientUtilTest`. `DynamoDbStreamsAsyncClientUtil` is gone, the source builds the http client and the async client itself, and its tests moved to `DynamoDbStreamsSourceAsyncClientTest`. ########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/DynamoDbStreamsSource.java: ########## @@ -259,7 +284,15 @@ private DynamoDbStreamsProxy createDynamoDbStreamsProxy( DynamodbStreamsSourceConfigConstants .BASE_DDB_STREAMS_USER_AGENT_PREFIX_FORMAT, DynamodbStreamsSourceConfigConstants.DDB_STREAMS_CLIENT_USER_AGENT_PREFIX); - return new DynamoDbStreamsProxy(dynamoDbStreamsClient, httpClient); + if (!withAsyncClient) { Review Comment: did both. in non-blocking mode the proxy now has only the async client and its http client, so no sync client and no second credentials provider. the proxy is also created with the first split reader (`LazyStreamProxy`), so a reader that never gets a split opens nothing. ########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/proxy/DynamoDbStreamsProxy.java: ########## @@ -276,4 +477,49 @@ private GetRecordsResponse getRecords(String shardIterator) { return dynamoDbStreamsClient.getRecords( GetRecordsRequest.builder().shardIterator(shardIterator).build()); } + + private CompletableFuture<GetRecordsResponse> getRecordsAsync( + DynamoDbStreamsAsyncClient asyncClient, @Nullable String shardIterator) { + if (Objects.isNull(shardIterator)) { + return CompletableFuture.completedFuture(EMPTY_GET_RECORDS_RESPONSE); + } + return asyncClient.getRecords( + GetRecordsRequest.builder().shardIterator(shardIterator).build()); + } + + private GetRecordsResponse storeNextShardIterator( + String shardId, GetRecordsResponse getRecordsResponse) { + if (getRecordsResponse.nextShardIterator() != null) { + shardIdToIteratorStore.put(shardId, getRecordsResponse.nextShardIterator()); + } + return getRecordsResponse; + } + + private DynamoDbStreamsAsyncClient getAsyncClient() { + if (dynamoDbStreamsAsyncClient == null) { + throw new IllegalStateException( + "DynamoDbStreamsProxy has no DynamoDbStreamsAsyncClient. " + + "Non-blocking calls are not available."); + } + return dynamoDbStreamsAsyncClient; + } + + private static void logIfStreamIsEnabling( Review Comment: done, both sync describeStream methods use `logIfStreamIsEnabling` now, and there's a `describeStreamRequest(streamArn, shardFilter)` helper. ########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/reader/PollingDynamoDbStreamsShardSplitReader.java: ########## @@ -144,14 +200,220 @@ public RecordsWithSplitIds<Record> fetch() throws IOException { splitContext.splitState.getStreamArn(), splitContext.splitState.getShardId(), splitContext.splitState.getNextStartingPosition()); + assignedSplits.add(splitContext); + return processGetRecordsResponse(splitContext, getRecordsResponse, currentTime); + } + + private RecordsWithSplitIds<Record> handleShardEnd( + DynamoDbStreamsShardSplitWithContext splitContext, long currentTime) { + if (!splitContext.hasAttemptedChildShardDiscovery) { + splitContext.hasAttemptedChildShardDiscovery = true; + splitContext.childShardDiscoveryAttempts = 0; + } + + if (splitContext.childShardDiscoveryAttempts < MAX_RETRY_ATTEMPTS_FOR_CHILD_SHARDS) { + long nextChildShardDiscoveryEligibleTime = + getNextEligibleTimeForChildDiscovery(splitContext); + if (currentTime >= nextChildShardDiscoveryEligibleTime) { + ListShardsResult listShardsResult = + dynamodbStreams.listShardsWithFilter( + splitContext.splitState.getStreamArn(), + childShardFilter(splitContext)); + RecordsWithSplitIds<Record> finishedSplit = + processChildShardDiscoveryResult( + splitContext, listShardsResult, currentTime); + if (finishedSplit != null) { + return finishedSplit; + } + } + assignedSplits.add(splitContext); + return INCOMPLETE_SHARD_EMPTY_RECORDS; + } else { + return finishedSplit(splitContext); + } + } + + private void sleep(long milliseconds) throws IOException { + try { + Thread.sleep(milliseconds); + } catch (InterruptedException e) { + throw new IOException("Split reader was interrupted while sleeping", e); + } + } + + // --------------------------------------------------------------------------------------- + // Non-blocking calls + // --------------------------------------------------------------------------------------- + + /** + * Starts a call for every shard that is ready and has no call in progress. Then collects one + * completed call. If no call is complete, waits until a call completes, a shard is ready, or + * {@link #wakeUp()} is called, for at most {@code NON_BLOCKING_FETCH_MAX_WAIT}. + */ + private RecordsWithSplitIds<Record> fetchNonBlocking() throws IOException { + long currentTime = System.currentTimeMillis(); + long nextEligibleTime = Long.MAX_VALUE; + + Iterator<DynamoDbStreamsShardSplitWithContext> splitIterator = assignedSplits.iterator(); + while (splitIterator.hasNext()) { + DynamoDbStreamsShardSplitWithContext splitContext = splitIterator.next(); + if (splitContext.callInProgress + || pausedSplitIds.contains(splitContext.splitState.getSplitId())) { + continue; + } + + long eligibleTime; + if (splitContext.splitState.isShardEndReached()) { + if (!splitContext.hasAttemptedChildShardDiscovery) { Review Comment: done, both paths use `isChildShardDiscoveryExhausted`, and `hasAttemptedChildShardDiscovery` is gone. ########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/reader/PollingDynamoDbStreamsShardSplitReader.java: ########## @@ -144,14 +200,220 @@ public RecordsWithSplitIds<Record> fetch() throws IOException { splitContext.splitState.getStreamArn(), splitContext.splitState.getShardId(), splitContext.splitState.getNextStartingPosition()); + assignedSplits.add(splitContext); + return processGetRecordsResponse(splitContext, getRecordsResponse, currentTime); + } + + private RecordsWithSplitIds<Record> handleShardEnd( + DynamoDbStreamsShardSplitWithContext splitContext, long currentTime) { + if (!splitContext.hasAttemptedChildShardDiscovery) { + splitContext.hasAttemptedChildShardDiscovery = true; + splitContext.childShardDiscoveryAttempts = 0; + } + + if (splitContext.childShardDiscoveryAttempts < MAX_RETRY_ATTEMPTS_FOR_CHILD_SHARDS) { + long nextChildShardDiscoveryEligibleTime = + getNextEligibleTimeForChildDiscovery(splitContext); + if (currentTime >= nextChildShardDiscoveryEligibleTime) { + ListShardsResult listShardsResult = + dynamodbStreams.listShardsWithFilter( + splitContext.splitState.getStreamArn(), + childShardFilter(splitContext)); + RecordsWithSplitIds<Record> finishedSplit = + processChildShardDiscoveryResult( + splitContext, listShardsResult, currentTime); + if (finishedSplit != null) { + return finishedSplit; + } + } + assignedSplits.add(splitContext); + return INCOMPLETE_SHARD_EMPTY_RECORDS; + } else { + return finishedSplit(splitContext); + } + } + + private void sleep(long milliseconds) throws IOException { + try { + Thread.sleep(milliseconds); + } catch (InterruptedException e) { + throw new IOException("Split reader was interrupted while sleeping", e); + } + } + + // --------------------------------------------------------------------------------------- + // Non-blocking calls + // --------------------------------------------------------------------------------------- + + /** + * Starts a call for every shard that is ready and has no call in progress. Then collects one + * completed call. If no call is complete, waits until a call completes, a shard is ready, or + * {@link #wakeUp()} is called, for at most {@code NON_BLOCKING_FETCH_MAX_WAIT}. + */ + private RecordsWithSplitIds<Record> fetchNonBlocking() throws IOException { + long currentTime = System.currentTimeMillis(); + long nextEligibleTime = Long.MAX_VALUE; + + Iterator<DynamoDbStreamsShardSplitWithContext> splitIterator = assignedSplits.iterator(); + while (splitIterator.hasNext()) { + DynamoDbStreamsShardSplitWithContext splitContext = splitIterator.next(); + if (splitContext.callInProgress + || pausedSplitIds.contains(splitContext.splitState.getSplitId())) { + continue; + } + + long eligibleTime; + if (splitContext.splitState.isShardEndReached()) { + if (!splitContext.hasAttemptedChildShardDiscovery) { + splitContext.hasAttemptedChildShardDiscovery = true; + splitContext.childShardDiscoveryAttempts = 0; + } + if (splitContext.childShardDiscoveryAttempts + >= MAX_RETRY_ATTEMPTS_FOR_CHILD_SHARDS) { + splitIterator.remove(); + return finishedSplit(splitContext); + } + eligibleTime = getNextEligibleTimeForChildDiscovery(splitContext); + if (currentTime >= eligibleTime) { + startChildShardDiscovery(splitContext, currentTime); + continue; + } + } else { + eligibleTime = getNextEligibleTime(splitContext); + if (currentTime >= eligibleTime) { + startGetRecords(splitContext, currentTime); + continue; + } + } + nextEligibleTime = Math.min(nextEligibleTime, eligibleTime); + } + + CompletedCall completedCall = + pollResponseQueue( + Math.min( + NON_BLOCKING_FETCH_MAX_WAIT.toMillis(), + nextEligibleTime - currentTime)); + if (completedCall == null || completedCall == WAKE_UP) { + return INCOMPLETE_SHARD_EMPTY_RECORDS; + } + return handleCompletedCall(completedCall); + } + + private void startGetRecords( + DynamoDbStreamsShardSplitWithContext splitContext, long startTime) { + LOG.debug( + "Starting GetRecords for split: {}, wasEmptyPoll: {}", + splitContext.splitState.getSplitId(), + splitContext.wasLastPollEmpty); + splitContext.callInProgress = true; + dynamodbStreams + .getRecordsAsync( + splitContext.splitState.getStreamArn(), + splitContext.splitState.getShardId(), + splitContext.splitState.getNextStartingPosition()) + .whenComplete( + (response, error) -> + responseQueue.add( + new CompletedCall( + splitContext, + startTime, + false, + response, + null, + error))); + } + + private void startChildShardDiscovery( + DynamoDbStreamsShardSplitWithContext splitContext, long startTime) { + LOG.debug( + "Starting child shard discovery for split: {}, attempt: {}", + splitContext.splitState.getSplitId(), + splitContext.childShardDiscoveryAttempts); + splitContext.callInProgress = true; + dynamodbStreams + .listShardsWithFilterAsync( + splitContext.splitState.getStreamArn(), childShardFilter(splitContext)) + .whenComplete( + (result, error) -> + responseQueue.add( + new CompletedCall( + splitContext, + startTime, + true, + null, + result, + error))); + } + + @Nullable + private CompletedCall pollResponseQueue(long waitMillis) throws IOException { + try { + return responseQueue.poll(waitMillis, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Split reader was interrupted while waiting for records", e); + } + } + + private RecordsWithSplitIds<Record> handleCompletedCall(CompletedCall completedCall) + throws IOException { + DynamoDbStreamsShardSplitWithContext splitContext = completedCall.splitContext; + if (completedCall.error != null) { + // The call failed after all SDK retries. Fail the fetch, as a blocking call does. + throw rethrow(completedCall.error); + } + if (pausedSplitIds.contains(splitContext.splitState.getSplitId())) { + // The call stays in progress until the split resumes, so that the reader does not + // start a second call for the split. + heldCalls.put(splitContext.splitState.getSplitId(), completedCall); + return INCOMPLETE_SHARD_EMPTY_RECORDS; + } + splitContext.callInProgress = false; + + if (!completedCall.childShardDiscovery) { + return processGetRecordsResponse( + splitContext, completedCall.getRecordsResponse, completedCall.startTime); + } + RecordsWithSplitIds<Record> finishedSplit = + processChildShardDiscoveryResult( + splitContext, completedCall.listShardsResult, completedCall.startTime); + if (finishedSplit == null) { + return INCOMPLETE_SHARD_EMPTY_RECORDS; + } + assignedSplits.remove(splitContext); + return finishedSplit; + } + + private static IOException rethrow(Throwable error) { Review Comment: done, using `stripCompletionException`, `rethrowIOException` and `findThrowable`. ########## flink-connector-aws/flink-connector-dynamodb/src/test/java/org/apache/flink/connector/dynamodb/source/reader/PollingDynamoDbStreamsShardSplitReaderTest.java: ########## @@ -69,14 +70,56 @@ public void init() { new DynamoDbStreamsShardMetrics( TestUtil.getTestSplit(TEST_SHARD_ID), metricListener.getMetricGroup())); splitReader = - new PollingDynamoDbStreamsShardSplitReader( + createSplitReader( testStreamProxy, NON_EMPTY_POLLING_DELAY_MILLIS, EMPTY_POLLING_DELAY_MILLIS, new ConcurrentHashMap<>(), shardMetricGroupMap); } + /** Returns true to run the tests with blocking GetRecords calls. */ + protected boolean blockingGetRecords() { + return false; Review Comment: moved it to the blocking subclass. `testWakeUpReturnsWaitingFetch` in the non-blocking test covers the wake-up in non-blocking mode. ########## flink-connector-aws/flink-connector-dynamodb/src/main/java/org/apache/flink/connector/dynamodb/source/reader/DynamoDbStreamsSourceReader.java: ########## @@ -68,7 +75,30 @@ public DynamoDbStreamsSourceReader( SourceReaderContext context, Map<String, List<Shard>> childShardIdMap, Map<String, DynamoDbStreamsShardMetrics> shardMetricGroupMap) { + this( + splitFetcherManager, + recordEmitter, + config, + context, + childShardIdMap, + shardMetricGroupMap, + null); Review Comment: done, removed all three, and the tests pass the arguments directly. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
