From bd0267e2923a19a247c5dee0745dbfd6e5b8bbe2 Mon Sep 17 00:00:00 2001
From: Stephen Margheim <stephen.margheim@gmail.com>
Date: Sun, 21 Jan 2024 18:41:55 +0100
Subject: [PATCH] Implement statement_timeout with timespecs

---
 ext/sqlite3/database.c     | 16 ++++++++++------
 ext/sqlite3/database.h     |  2 +-
 ext/sqlite3/sqlite3_ruby.h |  1 +
 ext/sqlite3/statement.c    |  2 +-
 ext/sqlite3/timespec.h     | 20 ++++++++++++++++++++
 5 files changed, 33 insertions(+), 8 deletions(-)
 create mode 100644 ext/sqlite3/timespec.h

diff --git a/ext/sqlite3/database.c b/ext/sqlite3/database.c
index ae00762b..7f48cef1 100644
--- a/ext/sqlite3/database.c
+++ b/ext/sqlite3/database.c
@@ -12,6 +12,11 @@
 
 VALUE cSqlite3Database;
 
+// Function to check if two timespec structures are equal
+int timespec_equal(const struct timespec* ts1, const struct timespec* ts2) {
+    return (ts1->tv_sec == ts2->tv_sec) && (ts1->tv_nsec == ts2->tv_nsec);
+}
+
 static void
 database_mark(void *ctx)
 {
@@ -266,12 +271,11 @@ rb_sqlite3_statement_timeout(void *context)
     struct timespec currentTime;
     clock_gettime(CLOCK_MONOTONIC, &currentTime);
 
-    if (ctx->stmt_deadline == 0) {
-        ctx->stmt_deadline = currentTime.tv_sec * 1e9 + currentTime.tv_nsec + (long long)ctx->stmt_timeout * 1e6;
-    } else {
-        long long timeDiff = ctx->stmt_deadline - (currentTime.tv_sec * 1e9 + currentTime.tv_nsec);
-
-        if (timeDiff < 0) { return 1; }
+    if (!timespecisset(&ctx->stmt_deadline)) {
+        // Set stmt_deadline if not already set
+        ctx->stmt_deadline = currentTime;
+    } else if (timespecafter(&currentTime, &ctx->stmt_deadline)) {
+        return 1;
     }
 
     return 0;
diff --git a/ext/sqlite3/database.h b/ext/sqlite3/database.h
index 9fbdb60d..3123f4fe 100644
--- a/ext/sqlite3/database.h
+++ b/ext/sqlite3/database.h
@@ -7,7 +7,7 @@ struct _sqlite3Ruby {
     sqlite3 *db;
     VALUE busy_handler;
     int stmt_timeout;
-    long long stmt_deadline;
+    struct timespec stmt_deadline;
 };
 
 typedef struct _sqlite3Ruby sqlite3Ruby;
diff --git a/ext/sqlite3/sqlite3_ruby.h b/ext/sqlite3/sqlite3_ruby.h
index bcf53e63..088d3cd5 100644
--- a/ext/sqlite3/sqlite3_ruby.h
+++ b/ext/sqlite3/sqlite3_ruby.h
@@ -41,6 +41,7 @@ extern VALUE cSqlite3Blob;
 #include <statement.h>
 #include <exception.h>
 #include <backup.h>
+#include <timespec.h>
 
 int bignum_to_int64(VALUE big, sqlite3_int64 *result);
 
diff --git a/ext/sqlite3/statement.c b/ext/sqlite3/statement.c
index 9367bc64..d049dec1 100644
--- a/ext/sqlite3/statement.c
+++ b/ext/sqlite3/statement.c
@@ -70,7 +70,7 @@ prepare(VALUE self, VALUE db, VALUE sql)
              );
 
     CHECK(db_ctx->db, status);
-    db_ctx->stmt_deadline = 0;
+    timespecclear(&db_ctx->stmt_deadline);
 
     return rb_str_new2(tail);
 }
diff --git a/ext/sqlite3/timespec.h b/ext/sqlite3/timespec.h
new file mode 100644
index 00000000..322fe758
--- /dev/null
+++ b/ext/sqlite3/timespec.h
@@ -0,0 +1,20 @@
+#define	timespecclear(tsp)		(tsp)->tv_sec = (tsp)->tv_nsec = 0
+#define	timespecisset(tsp)		((tsp)->tv_sec || (tsp)->tv_nsec)
+#define	timespecisvalid(tsp)						\
+	((tsp)->tv_nsec >= 0 && (tsp)->tv_nsec < 1000000000L)
+#define	timespeccmp(tsp, usp, cmp)					\
+	(((tsp)->tv_sec == (usp)->tv_sec) ?				\
+	    ((tsp)->tv_nsec cmp (usp)->tv_nsec) :			\
+	    ((tsp)->tv_sec cmp (usp)->tv_sec))
+#define	timespecsub(tsp, usp, vsp)					\
+	do {								\
+		(vsp)->tv_sec = (tsp)->tv_sec - (usp)->tv_sec;		\
+		(vsp)->tv_nsec = (tsp)->tv_nsec - (usp)->tv_nsec;	\
+		if ((vsp)->tv_nsec < 0) {				\
+			(vsp)->tv_sec--;				\
+			(vsp)->tv_nsec += 1000000000L;			\
+		}							\
+	} while (0)
+#define timespecafter(tsp, usp)						\
+	(((tsp)->tv_sec > (usp)->tv_sec) ||				\
+	 ((tsp)->tv_sec == (usp)->tv_sec && (tsp)->tv_nsec > (usp)->tv_nsec))