db: fix memory leaks

- sqlite3: table prevented pst gc unless close-on-exec
 - statement cache wrongly disabled close-on-exec
   (not a leak, just gc'd very slowly)
 - limit statement cache size

collects/db/private/generic/common.rkt

@@ -88,7 +88,6 @@
 
     (define/public (dprintf fmt . args)
       (when DEBUG? (apply fprintf (current-error-port) fmt args)))
-    
     ))
 
 ;; ----------------------------------------
@@ -422,23 +421,21 @@
              dprintf)
     (super-new)
 
+    (field [max-cache-size 50])
+
     ;; Statement Cache
     ;; updated by prepare; potentially invalidated by query (via check/invalidate-cache)
 
     (define pst-cache '#hash())
 
     (define/public (get-cached-statement stmt)
-      (cond [(use-cache?)
-             (let ([cached-pst (hash-ref pst-cache stmt #f)])
-               (cond [cached-pst
-                      (dprintf "  ** using cached statement\n")
-                      cached-pst]
-                     [else
-                      (dprintf "  ** statement not in cache\n")
-                      #f]))]
-            [else
-             (dprintf "  ** not using statement cache\n")
-             #f]))
+      (let ([cached-pst (hash-ref pst-cache stmt #f)])
+        (cond [cached-pst
+               (dprintf "  ** using cached statement\n")
+               cached-pst]
+              [else
+               (dprintf "  ** statement not in cache\n")
+               #f])))
 
     (define/public (safe-statement-type? stmt-type)
       (memq stmt-type '(select insert update delete with)))
@@ -456,25 +453,35 @@
         ((never) #f)
         ((in-transaction) (eq? (get-tx-status) #t))))
 
-    ;; check/invalidate-cache : statement/pst/symbol/#f -> hash/#f
+    ;; check/invalidate-cache : statement/pst -> hash/#f
     ;; Returns old cache on invalidation, or #f if stmt is safe.
+    ;; May also return part of old cache (excluding pst) when cache gets too big.
     (define/public (check/invalidate-cache x)
       #|
       Sufficient to check on every query execution whether statement type is safe
       (ie, SELECT, INSERT, etc). All statements sent as strings are considered
       unsafe, because they're usually transactional SQL.
       |#
+      (define (invalidate! except)
+        (dprintf "  ** invalidating statement cache~a\n" (if except " (too big)" ""))
+        (let ([cache pst-cache])
+          (set! pst-cache '#hash())
+          (cond [except
+                 (cache-statement! except)
+                 (hash-remove cache (send except get-stmt))]
+                [else
+                 cache])))
       (cond [(statement-binding? x)
              (check/invalidate-cache (statement-binding-pst x))]
             [(prepared-statement? x)
-             (check/invalidate-cache (send x get-stmt-type))]
-            [else
-             (cond [(safe-statement-type? x)
-                    #f]
-                   [else
-                    (dprintf "  ** invalidating statement cache\n")
-                    (begin0 pst-cache
-                      (set! pst-cache '#hash()))])]))
+             (let ([stmt-type (send x get-stmt-type)])
+               (cond [(safe-statement-type? stmt-type)
+                      (if (< (hash-count pst-cache) max-cache-size)
+                          #f
+                          (invalidate! x))]
+                     [else
+                      (invalidate! #f)]))]
+            [else (invalidate! #f)]))
 
     ;; Prepare
 
@@ -485,14 +492,16 @@
           (prepare1 fsym stmt close-on-exec?))))
 
     (define/public (prepare1 fsym stmt close-on-exec?)
-      (cond [close-on-exec?
+      (cond [(and close-on-exec? (use-cache?))
              (or (get-cached-statement stmt)
                  (let* ([stmt-type (classify-stmt stmt)]
                         [safe? (safe-statement-type? stmt-type)]
                         [pst (prepare1* fsym stmt (if safe? #f close-on-exec?) stmt-type)])
                    (when safe? (cache-statement! pst))
                    pst))]
-            [else (prepare1* fsym stmt #f (classify-stmt stmt))]))
+            [else
+             (dprintf "  ** not using statement cache\n")
+             (prepare1* fsym stmt close-on-exec? (classify-stmt stmt))]))
 
     (define/public (prepare1* fsym stmt close-on-exec?)
       (error/internal 'prepare1* "not implemented"))

collects/db/private/sqlite3/connection.rkt

@@ -20,7 +20,6 @@
                   busy-retry-delay)
 
     (define -db db)
-    (define statement-table (make-hasheq))
     (define saved-tx-status #f) ;; set by with-lock, only valid while locked
 
     (inherit call-with-lock*
@@ -190,21 +189,25 @@
                          (stmt-type stmt-type)
                          (stmt sql)
                          (owner this))])
-          (hash-set! statement-table pst #t)
           pst)))
 
     (define/override (disconnect* _politely?)
       (super disconnect* _politely?)
-      (start-atomic)
-      (let ([db -db])
-        (set! -db #f)
-        (end-atomic)
-        (when db
-          (let ([statements (hash-map statement-table (lambda (k v) k))])
-            (for ([pst (in-list statements)])
-              (do-free-statement 'disconnect pst))
-            (HANDLE 'disconnect (sqlite3_close db))
-            (void)))))
+      (call-as-atomic
+       (lambda ()
+         (let ([db -db])
+           (set! -db #f)
+           (when db
+             ;; Free all of connection's prepared statements. This will leave
+             ;; pst objects with dangling foreign objects, so don't try to free
+             ;; them again---check that -db is not-#f.
+             (let loop ()
+               (let ([stmt (sqlite3_next_stmt db #f)])
+                 (when stmt
+                   (HANDLE 'disconnect (sqlite3_finalize stmt))
+                   (loop))))
+             (HANDLE 'disconnect (sqlite3_close db))
+             (void))))))
 
     (define/public (get-base) this)
 
@@ -215,14 +218,13 @@
           (go)))
 
     (define/private (do-free-statement fsym pst)
-      (start-atomic)
-      (let ([stmt (send pst get-handle)])
-        (send pst set-handle #f)
-        (end-atomic)
-        (hash-remove! statement-table pst)
-        (when stmt
-          (HANDLE fsym (sqlite3_finalize stmt))
-          (void))))
+      (call-as-atomic
+       (lambda ()
+         (let ([stmt (send pst get-handle)])
+           (send pst set-handle #f)
+           (when (and stmt -db)
+             (HANDLE fsym (sqlite3_finalize stmt)))
+           (void)))))
 
     ;; Internal query
 

collects/tests/db/programs/memleak.rkt

@@ -0,0 +1,56 @@
+#lang racket/base
+(require (prefix-in db: db)
+         racket/match
+         racket/class)
+
+;; Test for db memory leaks
+
+;; FIXME: mysql quickly exhausts prepared statement limit
+
+(define the-dsn
+  (match (current-command-line-arguments)
+    [(vector x)
+     (string->symbol x)]
+    [_
+     (printf "No dsn argument given, exiting.\n")
+     (exit 0)]))
+
+;; 'tx, 'no
+(define start-tx-mode 'no)
+(define query-tx-mode 'no)
+
+;; boolean
+(define reconnect? #f)
+(define random-query? #f)
+
+;; ----
+
+(define (get-c)
+  (printf "-- connect\n")
+  (let ([c (db:dsn-connect the-dsn)])
+    (case start-tx-mode
+      ((tx) (db:start-transaction c))
+      ((no) (void)))
+    c))
+
+(define c #f)
+
+(define (c-test)
+  ;; Randomize to prevent statement caching
+  (define (go)
+    (db:query-value c (if random-query?
+                          (format "SELECT ~a" (random #e1e6))
+                          "SELECT 1")))
+  (case query-tx-mode
+    ((tx) (db:call-with-transaction c go))
+    ((no) (go))))
+
+(let loop ()
+  (collect-garbage)
+  (collect-garbage)
+  (when (or (not c) reconnect?)
+    (set! c (get-c)))
+  (displayln (quotient (current-memory-use) #e1e6))
+  (for ([i (in-range 10000)])
+    (c-test))
+  (loop))