From 77f1abf1e75cf25c033374f3b32c0b8435289ccf Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Wed, 26 Feb 2025 15:09:29 +0100 Subject: [PATCH] Don't crash on background worker connection error (#630) If you had an invalid MotherDuck token you would get into a crash-loop, because the background worker would fail to connect and crash Postgres. This fixes that by catching all exceptions thrown by C++ functions in the background worker. Resolves #622 --- src/pgduckdb_background_worker.cpp | 42 +++++++++++++++++------------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/src/pgduckdb_background_worker.cpp b/src/pgduckdb_background_worker.cpp index 01610ae..9810f56 100644 --- a/src/pgduckdb_background_worker.cpp +++ b/src/pgduckdb_background_worker.cpp @@ -73,6 +73,28 @@ typedef struct BackgoundWorkerShmemStruct { static BackgoundWorkerShmemStruct *BgwShmemStruct; +static duckdb::Connection * +BackgroundWorkerCheck(duckdb::Connection *connection, int64 *last_activity_count) { + if (!connection) { + connection = pgduckdb::DuckDBManager::Get().CreateConnection().get(); + } + SpinLockAcquire(&pgduckdb::BgwShmemStruct->lock); + int64 new_activity_count = pgduckdb::BgwShmemStruct->activity_count; + SpinLockRelease(&pgduckdb::BgwShmemStruct->lock); + if (*last_activity_count != new_activity_count) { + *last_activity_count = new_activity_count; + /* Trigger some activity to restart the syncing */ + pgduckdb::DuckDBQueryOrThrow(*connection, "FROM duckdb_tables() limit 0"); + } + /* + * If the extension is not registerid this loop is a no-op, which + * means we essentially keep polling until the extension is + * installed + */ + pgduckdb::SyncMotherDuckCatalogsWithPg(false, *connection->context); + return connection; +} + } // namespace pgduckdb extern "C" { @@ -92,7 +114,7 @@ pgduckdb_background_worker_main(Datum /* main_arg */) { pgduckdb::doing_motherduck_sync = true; pgduckdb::is_background_worker = true; - duckdb::unique_ptr connection; + duckdb::Connection *connection = nullptr; while (true) { // Initialize SPI (Server Programming Interface) and connect to the database @@ -102,23 +124,7 @@ pgduckdb_background_worker_main(Datum /* main_arg */) { PushActiveSnapshot(GetTransactionSnapshot()); if (pgduckdb::IsExtensionRegistered()) { - if (!connection) { - connection = pgduckdb::DuckDBManager::Get().CreateConnection(); - } - SpinLockAcquire(&pgduckdb::BgwShmemStruct->lock); - int64 new_activity_count = pgduckdb::BgwShmemStruct->activity_count; - SpinLockRelease(&pgduckdb::BgwShmemStruct->lock); - if (last_activity_count != new_activity_count) { - last_activity_count = new_activity_count; - /* Trigger some activity to restart the syncing */ - pgduckdb::DuckDBQueryOrThrow(*connection, "FROM duckdb_tables() limit 0"); - } - /* - * If the extension is not registerid this loop is a no-op, which - * means we essentially keep polling until the extension is - * installed - */ - pgduckdb::SyncMotherDuckCatalogsWithPg(false, *connection->context); + connection = InvokeCPPFunc(pgduckdb::BackgroundWorkerCheck, connection, &last_activity_count); } // Commit the transaction