Skip to content

Commit d78d875

Browse files
authored
std::atexit on parent
Exit if the parent thread exists
1 parent e0e16f2 commit d78d875

File tree

1 file changed

+28
-11
lines changed

1 file changed

+28
-11
lines changed

src/httpserver_extension.cpp

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@
77
#include "duckdb/main/extension_util.hpp"
88
#include "duckdb/common/atomic.hpp"
99
#include "duckdb/common/exception/http_exception.hpp"
10+
#include "duckdb/common/allocator.hpp"
1011

1112
#define CPPHTTPLIB_OPENSSL_SUPPORT
1213
#include "httplib.hpp"
1314

1415
#include <thread>
1516
#include <memory>
17+
#include <cstdlib>
1618

1719
namespace duckdb {
1820

@@ -21,7 +23,8 @@ struct HttpServerState {
2123
std::unique_ptr<std::thread> server_thread;
2224
std::atomic<bool> is_running;
2325
DatabaseInstance* db_instance;
24-
26+
unique_ptr<Allocator> allocator;
27+
2528
HttpServerState() : is_running(false), db_instance(nullptr) {}
2629
};
2730

@@ -32,16 +35,16 @@ static void HandleQuery(const string& query, duckdb_httplib_openssl::Response& r
3235
if (!global_state.db_instance) {
3336
throw IOException("Database instance not initialized");
3437
}
35-
38+
3639
Connection con(*global_state.db_instance);
3740
auto result = con.Query(query);
38-
41+
3942
if (result->HasError()) {
4043
res.status = 400;
4144
res.set_content(result->GetError(), "text/plain");
4245
return;
4346
}
44-
47+
4548
res.set_content(result->ToString(), "text/plain");
4649
} catch (const Exception& ex) {
4750
res.status = 400;
@@ -53,19 +56,22 @@ void HttpServerStart(DatabaseInstance& db, string_t host, int32_t port) {
5356
if (global_state.is_running) {
5457
throw IOException("HTTP server is already running");
5558
}
56-
59+
5760
global_state.db_instance = &db;
58-
global_state.server.reset(new duckdb_httplib_openssl::Server());
61+
global_state.server = make_uniq<duckdb_httplib_openssl::Server>();
5962
global_state.is_running = true;
60-
63+
64+
// Create a new allocator for the server thread
65+
global_state.allocator = make_uniq<Allocator>();
66+
6167
// Handle GET requests
6268
global_state.server->Get("/query", [](const duckdb_httplib_openssl::Request& req, duckdb_httplib_openssl::Response& res) {
6369
if (!req.has_param("q")) {
6470
res.status = 400;
6571
res.set_content("Missing query parameter 'q'", "text/plain");
6672
return;
6773
}
68-
74+
6975
auto query = req.get_param_value("q");
7076
HandleQuery(query, res);
7177
});
@@ -86,12 +92,12 @@ void HttpServerStart(DatabaseInstance& db, string_t host, int32_t port) {
8692
});
8793

8894
string host_str = host.GetString();
89-
global_state.server_thread.reset(new std::thread([host_str, port]() {
95+
global_state.server_thread = make_uniq<std::thread>([host_str, port]() {
9096
if (!global_state.server->listen(host_str.c_str(), port)) {
9197
global_state.is_running = false;
9298
throw IOException("Failed to start HTTP server on " + host_str + ":" + std::to_string(port));
9399
}
94-
}));
100+
});
95101
}
96102

97103
void HttpServerStop() {
@@ -104,17 +110,24 @@ void HttpServerStop() {
104110
global_state.server_thread.reset();
105111
global_state.db_instance = nullptr;
106112
global_state.is_running = false;
113+
114+
// Reset the allocator
115+
global_state.allocator.reset();
107116
}
108117
}
109118

119+
static void HttpServerCleanup() {
120+
HttpServerStop();
121+
}
122+
110123
static void LoadInternal(DatabaseInstance &instance) {
111124
auto httpserve_start = ScalarFunction("httpserve_start",
112125
{LogicalType::VARCHAR, LogicalType::INTEGER},
113126
LogicalType::VARCHAR,
114127
[&](DataChunk &args, ExpressionState &state, Vector &result) {
115128
auto &host_vector = args.data[0];
116129
auto &port_vector = args.data[1];
117-
130+
118131
UnaryExecutor::Execute<string_t, string_t>(
119132
host_vector, result, args.size(),
120133
[&](string_t host) {
@@ -134,6 +147,10 @@ static void LoadInternal(DatabaseInstance &instance) {
134147

135148
ExtensionUtil::RegisterFunction(instance, httpserve_start);
136149
ExtensionUtil::RegisterFunction(instance, httpserve_stop);
150+
151+
// Register the cleanup function to be called at exit
152+
std::atexit(HttpServerCleanup);
153+
137154
}
138155

139156
void HttpserverExtension::Load(DuckDB &db) {

0 commit comments

Comments
 (0)