Skip to content

Commit 161529a

Browse files
jayteejjayteej
committed
[FLINK-37949][Connectors/Kinesis] Fix client shutdown order to prevent Netty deadlocks
1 parent 373be6f commit 161529a

File tree

7 files changed

+151
-7
lines changed

7 files changed

+151
-7
lines changed

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/KinesisStreamsSource.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
import java.util.concurrent.ConcurrentHashMap;
8181
import java.util.function.Supplier;
8282

83+
import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.CLIENT_SHUTDOWN_TIMEOUT;
8384
import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_CONSUMER_NAME;
8485
import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_DESCRIBE_CONSUMER_RETRY_STRATEGY_MAX_ATTEMPTS_OPTION;
8586
import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_DESCRIBE_CONSUMER_RETRY_STRATEGY_MAX_DELAY_OPTION;
@@ -279,7 +280,10 @@ private KinesisAsyncStreamProxy createKinesisAsyncStreamProxy(
279280
KinesisAsyncClient.builder(),
280281
KinesisStreamsConfigConstants.BASE_KINESIS_USER_AGENT_PREFIX_FORMAT,
281282
KinesisStreamsConfigConstants.KINESIS_CLIENT_USER_AGENT_PREFIX);
282-
return new KinesisAsyncStreamProxy(kinesisAsyncClient, asyncHttpClient);
283+
return new KinesisAsyncStreamProxy(
284+
kinesisAsyncClient,
285+
asyncHttpClient,
286+
consumerConfig.get(CLIENT_SHUTDOWN_TIMEOUT).toMillis());
283287
}
284288

285289
private KinesisStreamProxy createKinesisStreamProxy(Configuration consumerConfig) {
@@ -319,7 +323,10 @@ private KinesisStreamProxy createKinesisStreamProxy(
319323
overrideBuilder,
320324
KinesisStreamsConfigConstants.BASE_KINESIS_USER_AGENT_PREFIX_FORMAT,
321325
KinesisStreamsConfigConstants.KINESIS_CLIENT_USER_AGENT_PREFIX);
322-
return new KinesisStreamProxy(kinesisClient, httpClient);
326+
return new KinesisStreamProxy(
327+
kinesisClient,
328+
httpClient,
329+
consumerConfig.get(CLIENT_SHUTDOWN_TIMEOUT).toMillis());
323330
}
324331

325332
private void setUpDeserializationSchema(SourceReaderContext sourceReaderContext)

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/config/KinesisSourceConfigOptions.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,4 +138,11 @@ public enum ConsumerLifecycle {
138138
.defaultValue(Duration.ofMillis(10000))
139139
.withDescription(
140140
"Timeout for consumer deregistration. When timeout is reached, code will continue as per normal.");
141+
142+
public static final ConfigOption<Duration> CLIENT_SHUTDOWN_TIMEOUT =
143+
ConfigOptions.key("source.client.shutdown.timeout")
144+
.durationType()
145+
.defaultValue(Duration.ofMillis(5000))
146+
.withDescription(
147+
"Timeout for graceful shutdown of the Kinesis clients.");
141148
}

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/proxy/KinesisAsyncStreamProxy.java‎

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,26 +21,41 @@
2121
import org.apache.flink.annotation.Internal;
2222
import org.apache.flink.connector.kinesis.source.split.StartingPosition;
2323

24+
import org.slf4j.Logger;
25+
import org.slf4j.LoggerFactory;
2426
import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
2527
import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
2628
import software.amazon.awssdk.services.kinesis.model.SubscribeToShardRequest;
2729
import software.amazon.awssdk.services.kinesis.model.SubscribeToShardResponseHandler;
2830

2931
import java.io.IOException;
32+
import java.time.Duration;
3033
import java.util.concurrent.CompletableFuture;
34+
import java.util.concurrent.TimeUnit;
35+
import java.util.concurrent.TimeoutException;
3136

3237
import static org.apache.flink.connector.kinesis.source.split.StartingPositionUtil.toSdkStartingPosition;
3338

3439
/** Implementation of async stream proxy for the Kinesis client. */
3540
@Internal
3641
public class KinesisAsyncStreamProxy implements AsyncStreamProxy {
42+
private static final Logger LOG = LoggerFactory.getLogger(KinesisAsyncStreamProxy.class);
43+
private static final long DEFAULT_SHUTDOWN_TIMEOUT_MILLIS = 5000;
44+
private final long shutdownTimeoutMillis;
45+
3746
private final KinesisAsyncClient kinesisAsyncClient;
3847
private final SdkAsyncHttpClient asyncHttpClient;
3948

4049
public KinesisAsyncStreamProxy(
4150
KinesisAsyncClient kinesisAsyncClient, SdkAsyncHttpClient asyncHttpClient) {
51+
this(kinesisAsyncClient, asyncHttpClient, DEFAULT_SHUTDOWN_TIMEOUT_MILLIS);
52+
}
53+
54+
public KinesisAsyncStreamProxy(
55+
KinesisAsyncClient kinesisAsyncClient, SdkAsyncHttpClient asyncHttpClient, long shutdownTimeoutMillis) {
4256
this.kinesisAsyncClient = kinesisAsyncClient;
4357
this.asyncHttpClient = asyncHttpClient;
58+
this.shutdownTimeoutMillis = shutdownTimeoutMillis;
4459
}
4560

4661
@Override
@@ -58,9 +73,49 @@ public CompletableFuture<Void> subscribeToShard(
5873
return kinesisAsyncClient.subscribeToShard(request, responseHandler);
5974
}
6075

76+
/**
77+
* Gracefully closes the Kinesis clients with a timeout.
78+
*
79+
* @param timeoutMillis maximum time to wait for clients to close
80+
*/
81+
public void gracefulClose(long timeoutMillis) {
82+
try {
83+
LOG.debug("Closing Kinesis clients with timeout of {} ms", timeoutMillis);
84+
85+
// Close the Kinesis client first with half the timeout
86+
long kinesisClientTimeout = timeoutMillis / 2;
87+
LOG.debug("Closing KinesisAsyncClient with timeout of {} ms", kinesisClientTimeout);
88+
try {
89+
CompletableFuture<Void> kinesisClientFuture = CompletableFuture.runAsync(() -> kinesisAsyncClient.close());
90+
kinesisClientFuture.get(kinesisClientTimeout, TimeUnit.MILLISECONDS);
91+
LOG.debug("Successfully closed KinesisAsyncClient");
92+
} catch (TimeoutException e) {
93+
LOG.warn("Closing KinesisAsyncClient timed out after {} ms", kinesisClientTimeout);
94+
} catch (Exception e) {
95+
LOG.warn("Error while closing KinesisAsyncClient", e);
96+
}
97+
98+
// Then close the HTTP client with the remaining timeout
99+
long httpClientTimeout = timeoutMillis - kinesisClientTimeout;
100+
LOG.debug("Closing SdkAsyncHttpClient with timeout of {} ms", httpClientTimeout);
101+
try {
102+
CompletableFuture<Void> httpClientFuture = CompletableFuture.runAsync(() -> asyncHttpClient.close());
103+
httpClientFuture.get(httpClientTimeout, TimeUnit.MILLISECONDS);
104+
LOG.debug("Successfully closed SdkAsyncHttpClient");
105+
} catch (TimeoutException e) {
106+
LOG.warn("Closing SdkAsyncHttpClient timed out after {} ms", httpClientTimeout);
107+
} catch (Exception e) {
108+
LOG.warn("Error while closing SdkAsyncHttpClient", e);
109+
}
110+
111+
LOG.debug("Completed graceful shutdown of Kinesis clients");
112+
} catch (Exception e) {
113+
LOG.warn("Error during graceful shutdown of Kinesis clients", e);
114+
}
115+
}
116+
61117
@Override
62118
public void close() throws IOException {
63-
kinesisAsyncClient.close();
64-
asyncHttpClient.close();
119+
gracefulClose(shutdownTimeoutMillis);
65120
}
66121
}

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/proxy/KinesisStreamProxy.java‎

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,21 +47,31 @@
4747
import java.util.ArrayList;
4848
import java.util.List;
4949
import java.util.Map;
50+
import java.util.concurrent.CompletableFuture;
5051
import java.util.concurrent.ConcurrentHashMap;
52+
import java.util.concurrent.TimeUnit;
53+
import java.util.concurrent.TimeoutException;
5154

5255
/** Implementation of the {@link StreamProxy} for Kinesis data streams. */
5356
@Internal
5457
public class KinesisStreamProxy implements StreamProxy {
5558
private static final Logger LOG = LoggerFactory.getLogger(KinesisStreamProxy.class);
59+
private static final long DEFAULT_SHUTDOWN_TIMEOUT_MILLIS = 5000;
60+
private final long shutdownTimeoutMillis;
5661

5762
private final KinesisClient kinesisClient;
5863
private final SdkHttpClient httpClient;
5964
private final Map<String, String> shardIdToIteratorStore;
6065

6166
public KinesisStreamProxy(KinesisClient kinesisClient, SdkHttpClient httpClient) {
67+
this(kinesisClient, httpClient, DEFAULT_SHUTDOWN_TIMEOUT_MILLIS);
68+
}
69+
70+
public KinesisStreamProxy(KinesisClient kinesisClient, SdkHttpClient httpClient, long shutdownTimeoutMillis) {
6271
this.kinesisClient = kinesisClient;
6372
this.httpClient = httpClient;
6473
this.shardIdToIteratorStore = new ConcurrentHashMap<>();
74+
this.shutdownTimeoutMillis = shutdownTimeoutMillis;
6575
}
6676

6777
@Override
@@ -201,9 +211,49 @@ public DescribeStreamConsumerResponse describeStreamConsumer(
201211
.build());
202212
}
203213

214+
/**
215+
* Gracefully closes the Kinesis clients with a timeout.
216+
*
217+
* @param timeoutMillis maximum time to wait for clients to close
218+
*/
219+
public void gracefulClose(long timeoutMillis) {
220+
try {
221+
LOG.debug("Closing Kinesis clients with timeout of {} ms", timeoutMillis);
222+
223+
// Close the Kinesis client first with half the timeout
224+
long kinesisClientTimeout = timeoutMillis / 2;
225+
LOG.debug("Closing KinesisClient with timeout of {} ms", kinesisClientTimeout);
226+
try {
227+
CompletableFuture<Void> kinesisClientFuture = CompletableFuture.runAsync(() -> kinesisClient.close());
228+
kinesisClientFuture.get(kinesisClientTimeout, TimeUnit.MILLISECONDS);
229+
LOG.debug("Successfully closed KinesisClient");
230+
} catch (TimeoutException e) {
231+
LOG.warn("Closing KinesisClient timed out after {} ms", kinesisClientTimeout);
232+
} catch (Exception e) {
233+
LOG.warn("Error while closing KinesisClient", e);
234+
}
235+
236+
// Then close the HTTP client with the remaining timeout
237+
long httpClientTimeout = timeoutMillis - kinesisClientTimeout;
238+
LOG.debug("Closing SdkHttpClient with timeout of {} ms", httpClientTimeout);
239+
try {
240+
CompletableFuture<Void> httpClientFuture = CompletableFuture.runAsync(() -> httpClient.close());
241+
httpClientFuture.get(httpClientTimeout, TimeUnit.MILLISECONDS);
242+
LOG.debug("Successfully closed SdkHttpClient");
243+
} catch (TimeoutException e) {
244+
LOG.warn("Closing SdkHttpClient timed out after {} ms", httpClientTimeout);
245+
} catch (Exception e) {
246+
LOG.warn("Error while closing SdkHttpClient", e);
247+
}
248+
249+
LOG.debug("Completed graceful shutdown of Kinesis clients");
250+
} catch (Exception e) {
251+
LOG.warn("Error during graceful shutdown of Kinesis clients", e);
252+
}
253+
}
254+
204255
@Override
205256
public void close() throws IOException {
206-
kinesisClient.close();
207-
httpClient.close();
257+
gracefulClose(shutdownTimeoutMillis);
208258
}
209259
}

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/fanout/FanOutKinesisShardSplitReader.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import org.apache.flink.connector.kinesis.source.split.KinesisShardSplit;
2828
import org.apache.flink.connector.kinesis.source.split.KinesisShardSplitState;
2929

30+
import org.slf4j.Logger;
31+
import org.slf4j.LoggerFactory;
3032
import software.amazon.awssdk.services.kinesis.model.SubscribeToShardEvent;
3133

3234
import java.time.Duration;
@@ -41,6 +43,8 @@
4143
*/
4244
@Internal
4345
public class FanOutKinesisShardSplitReader extends KinesisShardSplitReaderBase {
46+
private static final Logger LOG = LoggerFactory.getLogger(FanOutKinesisShardSplitReader.class);
47+
4448
private final AsyncStreamProxy asyncStreamProxy;
4549
private final String consumerArn;
4650
private final Duration subscriptionTimeout;
@@ -93,6 +97,18 @@ public void handleSplitsChanges(SplitsChange<KinesisShardSplit> splitsChanges) {
9397

9498
@Override
9599
public void close() throws Exception {
100+
LOG.debug("Closing FanOutKinesisShardSplitReader");
101+
102+
// Clear all subscriptions - they will be automatically cleaned up
103+
// when the AsyncStreamProxy is closed
104+
if (!splitSubscriptions.isEmpty()) {
105+
LOG.debug("Clearing {} active shard subscriptions", splitSubscriptions.size());
106+
splitSubscriptions.clear();
107+
}
108+
109+
// Close the proxy with a graceful timeout
110+
LOG.debug("Closing AsyncStreamProxy");
96111
asyncStreamProxy.close();
112+
LOG.debug("FanOutKinesisShardSplitReader closed");
97113
}
98114
}

‎flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/polling/PollingKinesisShardSplitReader.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.apache.flink.connector.kinesis.source.reader.KinesisShardSplitReaderBase;
2727
import org.apache.flink.connector.kinesis.source.split.KinesisShardSplitState;
2828

29+
import org.slf4j.Logger;
30+
import org.slf4j.LoggerFactory;
2931
import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse;
3032

3133
import java.util.Map;
@@ -36,6 +38,8 @@
3638
*/
3739
@Internal
3840
public class PollingKinesisShardSplitReader extends KinesisShardSplitReaderBase {
41+
private static final Logger LOG = LoggerFactory.getLogger(PollingKinesisShardSplitReader.class);
42+
3943
private final StreamProxy kinesis;
4044
private final Configuration configuration;
4145
private final int maxRecordsToGet;
@@ -65,6 +69,11 @@ protected RecordBatch fetchRecords(KinesisShardSplitState splitState) {
6569

6670
@Override
6771
public void close() throws Exception {
72+
LOG.debug("Closing PollingKinesisShardSplitReader");
73+
74+
// Close the proxy with a graceful timeout
75+
LOG.debug("Closing StreamProxy");
6876
kinesis.close();
77+
LOG.debug("PollingKinesisShardSplitReader closed");
6978
}
7079
}

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ under the License.
5454

5555
<properties>
5656
<aws.sdkv2.version>2.26.19</aws.sdkv2.version>
57-
<netty.version>4.1.86.Final</netty.version>
57+
<netty.version>4.1.122.Final</netty.version>
5858
<flink.version>2.0.0</flink.version>
5959
<jackson-bom.version>2.14.3</jackson-bom.version>
6060
<glue.schema.registry.version>1.1.18</glue.schema.registry.version>

0 commit comments

Comments
 (0)