[FLINK-13059][Cassandra Connector] Use getMaxConcurrentRequestsTimeou…

…t() also for flush
apache · Jul 6, 2019 · 7fe8595 · 7fe8595
1 parent 5f06981
commit 7fe8595
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/...ndra/src/main/java/org/apache/flink/streaming/connectors/cassandra/CassandraSinkBase.java b/...ndra/src/main/java/org/apache/flink/streaming/connectors/cassandra/CassandraSinkBase.java
@@ -164,8 +164,17 @@ private void checkAsyncErrors() throws Exception {
 		}
 	}
 
-	private void flush() {
-		semaphore.acquireUninterruptibly(config.getMaxConcurrentRequests());
+	private void flush() throws InterruptedException, TimeoutException {
+		if (!semaphore.tryAcquire(config.getMaxConcurrentRequests(), config.getMaxConcurrentRequestsTimeout().toMillis(), TimeUnit.MILLISECONDS)) {
+			throw new TimeoutException(
+				String.format(
+					"Failed to acquire all permits (%s) for flush in %s, only %s available",
+					config.getMaxConcurrentRequests(),
+					config.getMaxConcurrentRequestsTimeout(),
+					semaphore.availablePermits()
+				)
+			);
+		}
 		semaphore.release(config.getMaxConcurrentRequests());
 	}
 

diff --git a/.../src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraSinkBaseTest.java b/.../src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraSinkBaseTest.java
@@ -181,7 +181,7 @@ public void go() throws Exception {
 				}
 			};
 			t.start();
-			while (t.getState() != Thread.State.WAITING) {
+			while (t.getState() != Thread.State.TIMED_WAITING) {
 				Thread.sleep(5);
 			}
 
@@ -213,7 +213,7 @@ public void go() throws Exception {
 				}
 			};
 			t.start();
-			while (t.getState() != Thread.State.WAITING) {
+			while (t.getState() != Thread.State.TIMED_WAITING) {
 				Thread.sleep(5);
 			}