Fix transaction on_close and Java and Python block on close() #789

flyingsilverfin · 2025-09-26T21:35:51Z

Usage and product changes

We notice that calling transaction.close() does not wait until the server has freed up resource. This makes quick sequences, such as tests where transactions open and are followed by database deletes, unreliable. Further investigation that workarounds using the existing on_close callbacks in Python and Java caused segfaults. We fix both:

Transaction.close() in Python and Java now blocks for 1 round trip. We have the option to convert this fully into a promise return value the user can optionally resolve. We should pick the most relevant default!
We fix segfaults that occur when the Rust driver calls into Python/Java once the user attaches .on_close callbacks to transactions.

Implementation

Fix and enhance on_close callbacks:
- on attaching a callback, we don't return until the callback is actually registered (used to submit into an async channel, but not necessarily be recorded)
  - this is also sped up by having the lowest-level registration listener listen in an async context instead of a polling context
- we fix calling segfault that occurred on invoking the callback from Rust, mostly by enabling threading from the SWIG .i layer!
Make close() a blocking operation in Java and Python, which awaits a signal from the server that the transaction is actually closed and the resources are freed up.
We add on_close callback integration tests for Python, Java, and Rust

flyingsilverfin · 2025-09-30T17:33:38Z

c/typedb_driver.i

 * under the License.
 */

+%module(threads=1) typedb_driver


Key point! Since we're calling transaction close callbacks from another thread, we need this to avoid segfaults (not sure why - perhaps SWIG manages the GIL correctly for Python in this case!)

How did it work in 2.x? Or did it?..

flyingsilverfin · 2025-09-30T17:33:46Z

c/typedb_driver.i

-    virtual ~TransactionCallbackDirector() {}
-    virtual void callback(Error*) = 0;
+    virtual ~TransactionCallbackDirector() {
+        std::cout << "Callback::~Callback()" << std:: endl;


helpful in case things aren't working as expected

flyingsilverfin · 2025-09-30T17:34:14Z

python/tests/integration/BUILD

 )

+py_test(
+    name = "test_driver",


good place to leave future integration tests

c/typedb_driver.i

flyingsilverfin · 2025-09-30T18:27:14Z

java/connection/TransactionImpl.java

        if (nativeObject.isOwned()) {
            try {
-                transaction_force_close(nativeObject);
+                transaction_force_close(nativeObject).get();


Java: force resolve close(). We could also let the user resolve it themselves?

flyingsilverfin · 2025-09-30T18:27:22Z

java/test/integration/BUILD

 )

+typedb_java_test(
+    name = "test-driver",


new place to put java specific integration tests

flyingsilverfin · 2025-09-30T18:28:33Z

rust/src/connection/network/transmitter/transaction.rs

 impl Drop for TransactionTransmitter {
    fn drop(&mut self) {
-        self.force_close();
+        // TODO: in the async context, this now returns a promise... do we care?


Big open question @dmitrii-ubskii !

Ok: looks like the typical recommended thing to do is make it fire-and-forget. So basically users can call .force_close() if they want to await a promise first, then they can drop it if they want to be super sure in Rust.

Java/Python only use force_close() so they're blocking in GC anyway

flyingsilverfin · 2025-09-30T18:28:43Z

rust/src/connection/network/transmitter/transaction.rs


-    pub(in crate::connection) fn force_close(&self) {
+    #[cfg(not(feature = "sync"))]
+    pub(in crate::connection) fn force_close(&self) -> impl Promise<'static, Result<()>> {


Please validate these two functions!

flyingsilverfin · 2025-09-30T18:28:59Z

rust/src/connection/network/transmitter/transaction.rs

+        let (sender, mut sink) = unbounded_async();
+        self.on_close_register_sink.send((Box::new(callback), sender)).ok();
+        sink.blocking_recv().expect("Did not receive on_close registration success signal");


we now don't return to the user unti the on_close is actually registered

flyingsilverfin · 2025-09-30T18:29:46Z

rust/src/connection/network/transmitter/transaction.rs

+            move || Self::sync_dispatch_loop(queue_source, request_sink, collector, shutdown_signal)
        });
-        tokio::spawn(Self::listen_loop(response_source, collector, shutdown_sink));
+        tokio::spawn(Self::async_listen_loop(response_source, collector, on_close_callback_source, shutdown_sink));


Since we just have to get the on_close_callback into the Collector in some way, we can choose whether we lsiten to the channel in the async listen loop or the sync dispatch loop. To minimize waits, we're going to let Tokio manage it in the Async loop.

flyingsilverfin · 2025-09-30T18:30:48Z

rust/src/connection/transaction_stream.rs

+    pub(crate) fn force_close(&self)  -> impl Promise<'static, Result<()>> {
+        self.transaction_transmitter.force_close()


Note: in commit and rollback, we do something this:

let promise = self.single(TransactionRequest::Commit); promisify! { require_transaction_response!(resolve!(promise), Rollback)

Do we want to do anything like this here?

flyingsilverfin · 2025-09-30T18:34:18Z

rust/tests/integration/driver.rs

+            }
+        }));
+
+        drop(transaction); // TODO: drop isn't blocking... so we need to spin?


Drop question comes in here!

WIP Finally fixed on_close callback - threading issue Convert close to promise and resolve it in Python and Java Add Java onclose callback test Cleanup test Add mutex

farost · 2025-10-01T19:02:19Z

java/test/integration/BUILD

+        "@typedb_bazel_distribution//platform:is_linux_x86_64": "@typedb_artifact_linux-x86_64//file",
+        "@typedb_bazel_distribution//platform:is_mac_arm64": "@typedb_artifact_mac-arm64//file",
+        "@typedb_bazel_distribution//platform:is_mac_x86_64": "@typedb_artifact_mac-x86_64//file",
+#        "@typedb_bazel_distribution//platform:is_windows_x86_64": "@typedb_artifact_windows-x86_64//file",


Can be uncommented already

flyingsilverfin commented Sep 30, 2025

View reviewed changes

c/typedb_driver.i Outdated Show resolved Hide resolved

flyingsilverfin commented Sep 30, 2025

View reviewed changes

flyingsilverfin changed the title ~~Improve Rust Transaction on_close_callback~~ Fix transaction on_close and Java and Python block on close() Sep 30, 2025

flyingsilverfin marked this pull request as ready for review September 30, 2025 18:40

flyingsilverfin requested a review from dmitrii-ubskii September 30, 2025 18:40

flyingsilverfin added type: bug priority: high labels Sep 30, 2025

Improve Rust on_close_callback

e907782

WIP Finally fixed on_close callback - threading issue Convert close to promise and resolve it in Python and Java Add Java onclose callback test Cleanup test Add mutex

flyingsilverfin force-pushed the improve-on-close-callback branch from 8fb5f1e to e907782 Compare September 30, 2025 19:30

flyingsilverfin added 7 commits September 30, 2025 12:31

Trigger CI

a87a9d7

Deploy snapshot

df3ba3a

Fix build

b1b3f5b

Include mutex

ebf41f5

Fix drop behaviour

cbf1d84

Trigger CI

9e94c90

Use built-in Python 3.11 in mac jobs

a3bbf3a

farost reviewed Oct 1, 2025

View reviewed changes

		pub(crate) fn force_close(&self) -> impl Promise<'static, Result<()>> {
		self.transaction_transmitter.force_close()

Fix transaction on_close and Java and Python block on close() #789

Are you sure you want to change the base?

Fix transaction on_close and Java and Python block on close() #789

Uh oh!

Conversation

flyingsilverfin commented Sep 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Usage and product changes

Implementation

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

flyingsilverfin Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

flyingsilverfin commented Sep 26, 2025 •

edited

Loading

flyingsilverfin Sep 30, 2025 •

edited

Loading