Something went wrong. Try again.
Native PostgreSQL driver / client for Zig
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547const std = @import("std");const lib = @import("lib.zig");
const log = lib.log;const Conn = lib.Conn;const Result = lib.Result;const SSLCtx = lib.SSLCtx;const QueryRow = lib.QueryRow;const QueryRowUnsafe = lib.QueryRowUnsafe;const Listener = @import("listener.zig").Listener;
const Thread = std.Thread;const Allocator = std.mem.Allocator;
pub const Pool = struct { _opts: Opts, _timeout: u64, _conns: []*Conn, _available: usize, _missing: usize, _allocator: Allocator, _mutex: Thread.Mutex, _cond: Thread.Condition, _ssl_ctx: ?*lib.SSLCtx, _reconnector: Reconnector, _arena: std.heap.ArenaAllocator,
pub const Opts = struct { size: u16 = 10, auth: Conn.AuthOpts = .{}, connect: Conn.Opts = .{}, timeout: u32 = 10 * std.time.ms_per_s, connect_on_init_count: ?u16 = null, };
pub const Stats = struct { size: usize, available: usize, missing: usize, in_use: usize, };
pub fn initUri(allocator: Allocator, uri: std.Uri, opts: Opts) !*Pool { var po = try lib.parseOpts(uri, allocator); defer po.deinit(); po.opts.size = opts.size; po.opts.timeout = opts.timeout; return Pool.init(allocator, po.opts); }
pub fn init(allocator: Allocator, opts: Opts) !*Pool { var arena = std.heap.ArenaAllocator.init(allocator); const aa = arena.allocator(); errdefer arena.deinit();
const pool = try aa.create(Pool); const size = opts.size; const conns = try aa.alloc(*Conn, size);
var opts_copy = opts; var ssl_ctx: ?*SSLCtx = null; if (comptime lib.has_openssl) { switch (opts.connect.tls) { .off => {}, else => |tls_config| { if (opts.connect.host) |h| { opts_copy.connect._hostz = try aa.dupeZ(u8, h); } ssl_ctx = try lib.initializeSSLContext(tls_config); }, } } errdefer lib.freeSSLContext(ssl_ctx); const connect_on_init_count = opts.connect_on_init_count orelse size;
pool.* = .{ ._cond = .{}, ._mutex = .{}, ._conns = conns, ._arena = arena, ._opts = opts_copy, ._ssl_ctx = ssl_ctx, ._missing = 0, ._allocator = allocator, ._available = connect_on_init_count, ._reconnector = Reconnector.init(pool), ._timeout = @as(u64, @intCast(opts.timeout)) * std.time.ns_per_ms, };
var opened_connections: usize = 0; errdefer { for (0..opened_connections) |i| { conns[i].deinit(); } }
for (0..connect_on_init_count) |i| { conns[i] = try newConnection(pool, true); opened_connections += 1; }
const lazy_start_count = size - connect_on_init_count; pool._missing = lazy_start_count; for (0..lazy_start_count) |_| { try pool._reconnector.reconnect(); }
return pool; }
pub fn deinit(self: *Pool) void { self._reconnector.stop(); const allocator = self._allocator; for (self._conns) |conn| { conn.deinit(); allocator.destroy(conn); } lib.freeSSLContext(self._ssl_ctx); self._arena.deinit(); }
pub fn acquire(self: *Pool) !*Conn { const conns = self._conns; const deadline = std.time.nanoTimestamp() + @as(i64, @intCast(self._timeout));
self._mutex.lock(); errdefer self._mutex.unlock();
while (true) { const available = self._available; const missing = self._missing;
if (available == 0) { // Check if pool is completely exhausted const total_alive = self._conns.len - missing; if (total_alive == 0) { return error.PoolExhausted; }
lib.metrics.poolEmpty();
// Calculate remaining timeout const now = std.time.nanoTimestamp(); if (now >= deadline) { return error.Timeout; } const remaining_ns: u64 = @intCast(deadline - now);
try self._cond.timedWait(&self._mutex, remaining_ns); continue; }
const index = available - 1; const conn = conns[index]; self._available = index; self._mutex.unlock(); return conn; } }
pub fn release(self: *Pool, conn: *Conn) void { var conn_to_add = conn;
if (conn._state != .idle) { lib.metrics.poolDirty(); // conn should always be idle when being released. It's possible we can // recover from this (e.g. maybe we just need to read until we get a // ReadyForQuery), but we wouldn't want to block for too long. For now, // we'll just replace the connection. conn.deinit(); self._allocator.destroy(conn);
conn_to_add = newConnection(self, true) catch |err1| { // we failed to create the connection, track it as missing and let // the background reconnector try self._mutex.lock(); self._missing += 1; self._mutex.unlock();
self._reconnector.reconnect() catch |err2| { log.err("Re-opening connection failed ({}) and background reconnector failed to start ({})", .{ err1, err2 }); }; return; }; }
var conns = self._conns; self._mutex.lock(); const available = self._available; conns[available] = conn_to_add; self._available = available + 1; self._mutex.unlock(); self._cond.signal(); }
pub fn newListener(self: *Pool) !Listener { var listener = try Listener.open(self._allocator, self._opts.connect); try listener.auth(self._opts.auth); return listener; }
pub fn stats(self: *Pool) Stats { self._mutex.lock(); defer self._mutex.unlock();
const available = self._available; const missing = self._missing; const size = self._conns.len;
return .{ .size = size, .available = available, .missing = missing, .in_use = size - available - missing, }; }
pub fn exec(self: *Pool, sql: []const u8, values: anytype) !?i64 { return self.execOpts(sql, values, .{}); }
pub fn execOpts(self: *Pool, sql: []const u8, values: anytype, opts: Conn.QueryOpts) !?i64 { var conn = try self.acquire(); defer self.release(conn); return conn.execOpts(sql, values, opts); }
pub fn query(self: *Pool, sql: []const u8, values: anytype) !*Result { return self.queryOpts(sql, values, .{}); }
pub fn queryOpts(self: *Pool, sql: []const u8, values: anytype, opts_: Conn.QueryOpts) !*Result { var opts = opts_; opts.release_conn = true; var conn = try self.acquire(); errdefer self.release(conn); return conn.queryOpts(sql, values, opts); }
pub fn row(self: *Pool, sql: []const u8, values: anytype) !?QueryRow { return self.rowOpts(sql, values, .{}); }
pub fn rowUnsafe(self: *Pool, sql: []const u8, values: anytype) !?QueryRowUnsafe { return self.rowUnsafeOpts(sql, values, .{}); }
pub fn rowOpts(self: *Pool, sql: []const u8, values: anytype, opts_: Conn.QueryOpts) !?QueryRow { var opts = opts_; opts.release_conn = true; var conn = try self.acquire(); return conn.rowOpts(sql, values, opts); }
pub fn rowUnsafeOpts(self: *Pool, sql: []const u8, values: anytype, opts_: Conn.QueryOpts) !?QueryRowUnsafe { var opts = opts_; opts.release_conn = true; var conn = try self.acquire(); return conn.rowUnsafeOpts(sql, values, opts); }};
const Reconnector = struct { // number of connections that the pool is missing, i.e. how many need to be // reconnected count: usize,
// when stop is called, this is set to true stopped: bool,
pool: *Pool, mutex: Thread.Mutex,
// the thread, if any, that the monitor is running in thread: ?Thread,
fn init(pool: *Pool) Reconnector { return .{ .pool = pool, .count = 0, .mutex = .{}, .stopped = false, .thread = null, }; }
fn run(self: *Reconnector) void { const pool = self.pool; const retry_delay = 2 * std.time.ns_per_s;
self.mutex.lock(); defer self.mutex.unlock(); loop: while (self.count > 0) { const stopped = self.stopped; self.mutex.unlock(); if (stopped == true) { return; }
const conn = newConnection(pool, false) catch { std.Thread.sleep(retry_delay); self.mutex.lock(); continue :loop; };
// Decrement missing count when successfully recreated pool._mutex.lock(); std.debug.assert(pool._missing > 0); pool._missing -= 1; pool._mutex.unlock();
conn.release(); // inserts it into the pool self.mutex.lock(); self.count -= 1; }
self.thread.?.detach(); self.thread = null; }
fn stop(self: *Reconnector) void { self.mutex.lock(); self.stopped = true; self.mutex.unlock(); if (self.thread) |thrd| { thrd.join(); } }
fn reconnect(self: *Reconnector) !void { self.mutex.lock(); defer self.mutex.unlock(); self.count += 1; if (self.thread == null) { self.thread = try Thread.spawn(.{ .stack_size = 1024 * 1024 }, Reconnector.run, .{self}); } }};
fn newConnection(pool: *Pool, log_failure: bool) !*Conn { const opts = &pool._opts; const allocator = pool._allocator;
const conn = allocator.create(Conn) catch |err| { if (log_failure) log.err("connect error: {}", .{err}); return err; }; errdefer allocator.destroy(conn);
conn.* = Conn.open(allocator, opts.connect) catch |err| { if (log_failure) log.err("connect error: {}", .{err}); return err; }; errdefer conn.deinit();
conn.auth(opts.auth) catch |err| { if (log_failure) { if (conn.err) |pg_err| { log.err("connect error: {s}", .{pg_err.message}); } else { log.err("connect error: {}", .{err}); } } return err; }; conn._pool = pool; return conn;}
const t = lib.testing;test "Pool" { var pool = try Pool.init(t.allocator, .{ .size = 2, .auth = t.authOpts(.{}), .connect_on_init_count = 1, }); defer pool.deinit();
{ const c1 = try pool.acquire(); defer pool.release(c1); _ = try c1.exec( \\ drop table if exists pool_test; \\ create table pool_test (id int not null) , .{}); }
const t1 = try std.Thread.spawn(.{}, testPool, .{pool}); const t2 = try std.Thread.spawn(.{}, testPool, .{pool}); const t3 = try std.Thread.spawn(.{}, testPool, .{pool});
t1.join(); t2.join(); t3.join();
{ const c1 = try pool.acquire(); defer c1.release();
const affected = try c1.exec("delete from pool_test", .{}); try t.expectEqual(1500, affected.?); }}
test "Pool: Release" { var pool = try Pool.init(t.allocator, .{ .size = 2, .auth = .{ .database = "postgres", .username = "postgres", .password = "postgres", }, }); defer pool.deinit();
const c1 = try pool.acquire(); c1._state = .query; pool.release(c1);}
test "Pool: stats" { var pool = try Pool.init(t.allocator, .{ .size = 3, .auth = t.authOpts(.{}), }); defer pool.deinit();
// Initial state: all connections available { const s = pool.stats(); try t.expectEqual(3, s.size); try t.expectEqual(3, s.available); try t.expectEqual(0, s.missing); try t.expectEqual(0, s.in_use); }
// Acquire one connection const c1 = try pool.acquire(); { const s = pool.stats(); try t.expectEqual(3, s.size); try t.expectEqual(2, s.available); try t.expectEqual(0, s.missing); try t.expectEqual(1, s.in_use); }
// Acquire another const c2 = try pool.acquire(); { const s = pool.stats(); try t.expectEqual(3, s.size); try t.expectEqual(1, s.available); try t.expectEqual(0, s.missing); try t.expectEqual(2, s.in_use); }
// Release one pool.release(c1); { const s = pool.stats(); try t.expectEqual(3, s.size); try t.expectEqual(2, s.available); try t.expectEqual(0, s.missing); try t.expectEqual(1, s.in_use); }
// Release the other pool.release(c2); { const s = pool.stats(); try t.expectEqual(3, s.size); try t.expectEqual(3, s.available); try t.expectEqual(0, s.missing); try t.expectEqual(0, s.in_use); }}
test "Pool: exec" { var pool = try Pool.init(t.allocator, .{ .size = 1, .auth = t.authOpts(.{}) }); defer pool.deinit();
{ const n = try pool.exec("insert into simple_table values ($1), ($2), ($3)", .{ "pool_insert_args_a", "pool_insert_args_b", "pool_insert_args_c" }); try t.expectEqual(3, n.?); }
{ // this makes sure the connection was returned to the pool const n = try pool.exec("insert into simple_table values ($1)", .{"pool_insert_args_a"}); try t.expectEqual(1, n.?); }}
test "Pool: Query/Row" { var pool = try Pool.init(t.allocator, .{ .size = 1, .auth = t.authOpts(.{}) }); defer pool.deinit();
{ _ = try pool.exec("insert into all_types (id, col_int8, col_text) values ($1, $2, $3)", .{ 100, 1, "val-1" }); _ = try pool.exec("insert into all_types (id, col_int8, col_text) values ($1, $2, $3)", .{ 101, 2, "val-2" }); }
for (0..3) |_| { var result = try pool.query("select col_int8, col_text from all_types where id = any($1)", .{[2]i32{ 100, 101 }}); defer result.deinit();
const row1 = (try result.nextUnsafe()) orelse unreachable; try t.expectEqual(1, row1.get(i64, 0)); try t.expectString("val-1", row1.get([]u8, 1));
const row2 = (try result.nextUnsafe()) orelse unreachable; try t.expectEqual(2, row2.get(i64, 0)); try t.expectString("val-2", row2.get([]u8, 1));
try t.expectEqual(null, result.nextUnsafe()); }
for (0..3) |_| { var row = try pool.rowUnsafe("select col_int8, col_text from all_types where id = $1", .{101}) orelse unreachable; defer row.deinit() catch {};
try t.expectEqual(2, row.get(i64, 0)); try t.expectString("val-2", row.get([]u8, 1)); }}
test "Pool: Row error" { var pool = try Pool.init(t.allocator, .{ .size = 1, .auth = t.authOpts(.{}) }); defer pool.deinit();
_ = try pool.rowUnsafe("insert into all_types (id) values ($1)", .{200});
// This would segfault: // https://github.com/karlseguin/pg.zig/issues/34 try t.expectError(error.PG, pool.rowUnsafe("insert into all_types (id) values ($1)", .{200}));
try t.expectEqual(1, pool._available);}
fn testPool(p: *Pool) void { for (0..500) |i| { const conn = p.acquire() catch unreachable; _ = conn.exec("insert into pool_test (id) values ($1)", .{i}) catch unreachable; conn.release(); }}