From 552fb62f9dbc6d70fc027494d3665097170837cd Mon Sep 17 00:00:00 2001
From: Eitaro Fukamachi <e.arrows@gmail.com>
Date: Sat, 7 Sep 2024 13:29:57 +0000
Subject: [PATCH] Add `disconnect-cached-all`.

---
 dbi.asd               |  3 ++-
 src/cache/single.lisp | 13 ++++++++++++-
 src/cache/thread.lisp |  5 +++--
 src/dbi.lisp          |  6 ++++++
 t/dbi.lisp            | 37 +++++++++++++++++++++++++++++++++++++
 5 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 t/dbi.lisp

diff --git a/dbi.asd b/dbi.asd
index 63dfa46..146b23b 100644
--- a/dbi.asd
+++ b/dbi.asd
@@ -51,5 +51,6 @@
                   :components
                   ((:file "sqlite3")
                    (:file "postgres")
-                   (:file "mysql"))))))
+                   (:file "mysql")))
+                 (:file "dbi"))))
   :perform (test-op (op c) (symbol-call '#:rove '#:run c)))
diff --git a/src/cache/single.lisp b/src/cache/single.lisp
index 9b2a785..11e21fa 100644
--- a/src/cache/single.lisp
+++ b/src/cache/single.lisp
@@ -29,6 +29,17 @@
     (setf (gethash key cache) object)))
 
 ;; Just do nothing since it's single-threaded and the thread is obviously alive.
-(defun cleanup-cache-pool (pool)
+(defun cleanup-cache-pool (pool &key force)
   (declare (ignore pool))
+  (when force
+    (let ((cache (cache-pool-cache pool))
+          (cleanup-fn (cache-pool-cleanup-fn pool)))
+      (when cleanup-fn
+        (maphash (lambda (key conn)
+                   (declare (ignore key))
+                   (when conn
+                     (funcall cleanup-fn conn)))
+                 cache))
+      (setf (cache-pool-cache pool)
+            (make-hash-table :test 'equal))))
   (values))
diff --git a/src/cache/thread.lisp b/src/cache/thread.lisp
index 21f7dee..e1f6811 100644
--- a/src/cache/thread.lisp
+++ b/src/cache/thread.lisp
@@ -33,11 +33,12 @@
       (funcall (cache-pool-cleanup-fn pool) old-object))
     (setf (gethash key cache) object)))
 
-(defun cleanup-cache-pool (pool)
+(defun cleanup-cache-pool (pool &key force)
   (let ((cleanup-fn (cache-pool-cleanup-fn pool)))
     (bt2:with-lock-held ((cache-pool-lock pool))
       (maphash (lambda (thread cache)
-                 (unless (bt2:thread-alive-p thread)
+                 (when (or force
+                           (not (bt2:thread-alive-p thread)))
                    (when cleanup-fn
                      (maphash (lambda (_ conn)
                                 (declare (ignore _))
diff --git a/src/dbi.lisp b/src/dbi.lisp
index f528925..e72b2d5 100644
--- a/src/dbi.lisp
+++ b/src/dbi.lisp
@@ -7,6 +7,7 @@
                  #:dbi.logger)
   (:export #:connect
            #:connect-cached
+           #:disconnect-cached-all
            #:with-connection))
 (in-package #:dbi)
 
@@ -47,6 +48,11 @@
                   (apply #'connect connect-args))
           (cleanup-cache-pool pool)))))
 
+(defun disconnect-cached-all ()
+  (let ((pool *threads-connection-pool*))
+    (cleanup-cache-pool pool :force t))
+  (values))
+
 (defmacro with-autoload-on-missing (&body body)
   (let ((retrying (gensym))
         (e (gensym)))
diff --git a/t/dbi.lisp b/t/dbi.lisp
new file mode 100644
index 0000000..72c0624
--- /dev/null
+++ b/t/dbi.lisp
@@ -0,0 +1,37 @@
+(defpackage #:dbi-test.dbi
+  (:use #:cl
+        #:dbi
+        #:rove))
+(in-package #:dbi-test.dbi)
+
+(deftest disconnect-cached-all
+  (ok (= 0
+         (hash-table-count (dbi.cache.thread::cache-pool-hash dbi::*threads-connection-pool*))))
+  (let ((threads
+          (loop repeat 4
+                collect
+                   (bt2:make-thread
+                    (lambda ()
+                      (let ((conn
+                              (dbi:connect-cached
+                               :postgres
+                               :database-name (or (uiop:getenv "POSTGRES_DBNAME")
+                                                  "cl-dbi")
+                               :host (or (uiop:getenv "POSTGRES_HOST")
+                                         "localhost")
+                               :port (parse-integer
+                                      (or (uiop:getenv "POSTGRES_PORT")
+                                          "5432"))
+                               :username (or (uiop:getenv "POSTGRES_USER")
+                                             "nobody")
+                               :password (or (uiop:getenv "POSTGRES_PASS")
+                                             "nobody"))))
+                        (loop
+                          (sleep 1))))))))
+    (sleep 1)
+    (ok (= 4
+           (hash-table-count (dbi.cache.thread::cache-pool-hash dbi::*threads-connection-pool*))))
+    (dbi:disconnect-cached-all)
+    (ok (= 0
+           (hash-table-count (dbi.cache.thread::cache-pool-hash dbi::*threads-connection-pool*))))
+    (mapc #'bt2:destroy-thread threads)))