Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JAVA-3166: Enhanced exception handling of C* driver executeAsync() #1370

Merged
merged 1 commit into from
Aug 27, 2024
Merged

JAVA-3166: Enhanced exception handling of C* driver executeAsync() #1370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions connector/src/main/scala/com/datastax/spark/connector/writer/AsyncExecutor.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
package com.datastax.spark.connector.writer

import java.util.concurrent.{CompletionStage, Semaphore}
import java.util.concurrent.{CompletableFuture, CompletionStage, Semaphore}
import java.util.function.BiConsumer

import com.datastax.spark.connector.util.Logging

import scala.jdk.CollectionConverters._
Expand Down Expand Up @@ -41,9 +40,14 @@ class AsyncExecutor[T, R](asyncAction: T => CompletionStage[R], maxConcurrentTas
val executionTimestamp = System.nanoTime()

def tryFuture(): Future[R] = {
val value = asyncAction(task)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe something more Scala'ish - Try { ... } map { ... } ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed.


value.whenComplete(new BiConsumer[R, Throwable] {
val value = Try(asyncAction(task)) recover {
case e =>
val future = new CompletableFuture[R]()
future.completeExceptionally(e)
future
}

value.get.whenComplete(new BiConsumer[R, Throwable] {
private def release() {
semaphore.release()
pendingFutures.remove(promise.future)
Expand Down
18 changes: 18 additions & 0 deletions connector/src/test/scala/com/datastax/spark/connector/writer/AsyncExecutorTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package com.datastax.spark.connector.writer

import com.datastax.oss.driver.api.core.cql.{AsyncResultSet, SimpleStatement, Statement}

import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.{Callable, CompletableFuture, CompletionStage}

Expand Down Expand Up @@ -55,4 +57,20 @@ class AsyncExecutorTest {
totalFinishedExecutionsCounter.get() shouldBe taskCount
asyncExecutor.getLatestException() shouldBe None
}

@Test
def testGracefullyHandleCqlSessionExecuteExceptions() {
val executor = new AsyncExecutor[Statement[_], AsyncResultSet](
_ => {
// simulate exception returned by session.executeAsync() (not future)
throw new IllegalStateException("something bad happened")
}, 10, None, None
)
val stmt = SimpleStatement.newInstance("INSERT INTO table1 (key, value) VALUES (1, '100')");
val future = executor.executeAsync(stmt)
assertTrue(future.isCompleted)
val value = future.value.get
assertTrue(value.isInstanceOf[Failure[_]])
assertTrue(value.asInstanceOf[Failure[_]].exception.isInstanceOf[IllegalStateException])
}
}
Loading