Add a new RegisterBusyHandler method to SQLiteConn

Author: polyscone

_example/busy_handler/main.go

@@ -0,0 +1,80 @@
+package main
+
+import (
+	"database/sql"
+	"flag"
+	"time"
+
+	"github.com/mattn/go-sqlite3"
+)
+
+// A busy handler can be as simple as a hardcoded sleep
+// We use the count along with the sleep duration to decide
+// roughly when we've hit the timeout
+func simpleBusyCallback(count int) int {
+	const timeout = 5 * time.Second
+	const delay = 1 * time.Millisecond
+
+	if delay*time.Duration(count) >= timeout {
+		// Trigger a SQLITE_BUSY error
+		return 0
+	}
+
+	time.Sleep(delay)
+
+	// Attempt to access the database again
+	return 1
+}
+
+// This is a copy of the sqliteDefaultBusyCallback implementation
+// from the SQLite3 source code with minor changes
+//
+// This is equivalent to the function that's used when the
+// busy_timeout pragma is set
+func defaultBusyCallback(count int) int {
+	// All durations are in milliseconds
+	const timeout = 5000
+	delays := [...]int{1, 2, 5, 10, 15, 20, 25, 25, 25, 50, 50, 100}
+	totals := [...]int{0, 1, 3, 8, 18, 33, 53, 78, 103, 128, 178, 228}
+
+	var delay, prior int
+	if count < len(delays) {
+		delay = delays[count]
+		prior = totals[count]
+	} else {
+		delay = delays[len(delays)-1]
+		prior = totals[len(totals)-1] + delay*(count-len(totals))
+	}
+
+	if prior+delay > timeout {
+		delay = timeout - prior
+
+		if delay <= 0 {
+			// Trigger a SQLITE_BUSY error
+			return 0
+		}
+	}
+
+	time.Sleep(time.Duration(delay) * time.Millisecond)
+
+	// Attempt to access the database again
+	return 1
+}
+
+func main() {
+	var simple bool
+	flag.BoolVar(&simple, "simple", false, "Use the simple busy handler")
+	flag.Parse()
+
+	sql.Register("sqlite3_with_busy_handler_example", &sqlite3.SQLiteDriver{
+		ConnectHook: func(conn *sqlite3.SQLiteConn) error {
+			if simple {
+				conn.RegisterBusyHandler(simpleBusyCallback)
+			} else {
+				conn.RegisterBusyHandler(defaultBusyCallback)
+			}
+
+			return nil
+		},
+	})
+}

callback.go

@@ -97,6 +97,12 @@ func preUpdateHookTrampoline(handle unsafe.Pointer, dbHandle uintptr, op int, db
 	callback(data)
 }
 
+//export busyHandlerTrampoline
+func busyHandlerTrampoline(handle unsafe.Pointer, count int) int {
+	callback := lookupHandle(handle).(func(int) int)
+	return callback(count)
+}
+
 // Use handles to avoid passing Go pointers to C.
 type handleVal struct {
 	db  *SQLiteConn

sqlite3.go

@@ -170,6 +170,8 @@ void updateHookTrampoline(void*, int, char*, char*, sqlite3_int64);
 
 int authorizerTrampoline(void*, int, char*, char*, char*, char*);
 
+int busyHandlerTrampoline(void*, int);
+
 #ifdef SQLITE_LIMIT_WORKER_THREADS
 # define _SQLITE_HAS_LIMIT
 # define SQLITE_LIMIT_LENGTH                    0
@@ -684,6 +686,34 @@ func sqlite3CreateFunction(db *C.sqlite3, zFunctionName *C.char, nArg C.int, eTe
 	return C._sqlite3_create_function(db, zFunctionName, nArg, eTextRep, C.uintptr_t(uintptr(pApp)), (*[0]byte)(xFunc), (*[0]byte)(xStep), (*[0]byte)(xFinal))
 }
 
+// RegisterBusyHandler sets the busy handler for a connection.
+//
+// The parameter to the callback is the number of times that the busy
+// handler has been invoked previously for the same locking event.
+//
+// If there is an existing busy handler for this connection, it will be
+// removed. If callback is nil the existing handler (if any) will be removed
+// without creating a new one.
+//
+// If the busy callback returns 0 then no additional attempts are made to
+// access the database and SQLITE_BUSY is returned to the application.
+// If the callback returns non-zero, then another attempt is made to access
+// the database.
+func (c *SQLiteConn) RegisterBusyHandler(callback func(int) int) error {
+	var rv C.int
+	if callback == nil {
+		rv = C.sqlite3_busy_handler(c.db, nil, nil)
+	} else {
+		handle := newHandle(c, callback)
+		rv = C.sqlite3_busy_handler(c.db, (*[0]byte)(unsafe.Pointer(C.busyHandlerTrampoline)), handle)
+	}
+	if rv != C.SQLITE_OK {
+		return c.lastError()
+	}
+
+	return nil
+}
+
 // RegisterAggregator makes a Go type available as a SQLite aggregation function.
 //
 // Because aggregation is incremental, it's implemented in Go with a