Add support for sqlite3_unlock_notify

Mura Li · typeless · commit 0fce88d511f5 · 2018-10-01T15:26:06.000+08:00
diff --git a/.travis.yml b/.travis.yml
@@ -12,7 +12,7 @@ env:
   matrix:
     - GOTAGS=
     - GOTAGS=libsqlite3
-    - GOTAGS="sqlite_allow_uri_authority sqlite_app_armor sqlite_foreign_keys sqlite_fts5 sqlite_icu sqlite_introspect sqlite_json sqlite_secure_delete sqlite_see sqlite_stat4 sqlite_trace sqlite_userauth sqlite_vacuum_incr sqlite_vtable"
+    - GOTAGS="sqlite_allow_uri_authority sqlite_app_armor sqlite_foreign_keys sqlite_fts5 sqlite_icu sqlite_introspect sqlite_json sqlite_secure_delete sqlite_see sqlite_stat4 sqlite_trace sqlite_userauth sqlite_vacuum_incr sqlite_vtable sqlite_unlock_notify"
     - GOTAGS=sqlite_vacuum_full
 
 go:
diff --git a/sqlite3.go b/sqlite3.go
@@ -78,8 +78,38 @@ _sqlite3_exec(sqlite3* db, const char* pcmd, long long* rowid, long long* change
   return rv;
 }
 
+#ifdef SQLITE_ENABLE_UNLOCK_NOTIFY
+extern int sqlite3_step_blocking(sqlite3_stmt *stmt);
+extern int _sqlite3_step_blocking(sqlite3_stmt* stmt, long long* rowid, long long* changes);
+extern int sqlite3_prepare_v2_blocking(sqlite3 *db, const char *zSql, int nBytes, sqlite3_stmt **ppStmt, const char **pzTail);
+
+static int
+sqlite3_step_internal(sqlite3_stmt *stmt)
+{
+  return sqlite3_step_blocking(stmt);
+}
+
+static int
+_sqlite3_step_internal(sqlite3_stmt* stmt, long long* rowid, long long* changes)
+{
+  return _sqlite3_step_blocking(stmt, rowid, changes);
+}
+
 static int
-_sqlite3_step(sqlite3_stmt* stmt, long long* rowid, long long* changes)
+sqlite3_prepare_v2_internal(sqlite3 *db, const char *zSql, int nBytes, sqlite3_stmt **ppStmt, const char **pzTail)
+{
+  return sqlite3_prepare_v2_blocking(db, zSql, nBytes, ppStmt, pzTail);
+}
+
+#else
+static int
+sqlite3_step_internal(sqlite3_stmt *stmt)
+{
+  return sqlite3_step(stmt);
+}
+
+static int
+_sqlite3_step_internal(sqlite3_stmt* stmt, long long* rowid, long long* changes)
 {
   int rv = sqlite3_step(stmt);
   sqlite3* db = sqlite3_db_handle(stmt);
@@ -88,6 +118,13 @@ _sqlite3_step(sqlite3_stmt* stmt, long long* rowid, long long* changes)
   return rv;
 }
 
+static int
+sqlite3_prepare_v2_internal(sqlite3 *db, const char *zSql, int nBytes, sqlite3_stmt **ppStmt, const char **pzTail)
+{
+  return sqlite3_prepare_v2(db, zSql, nBytes, ppStmt, pzTail);
+}
+#endif
+
 void _sqlite3_result_text(sqlite3_context* ctx, const char* s) {
   sqlite3_result_text(ctx, s, -1, &free);
 }
@@ -1637,7 +1674,7 @@ func (c *SQLiteConn) prepare(ctx context.Context, query string) (driver.Stmt, er
 	defer C.free(unsafe.Pointer(pquery))
 	var s *C.sqlite3_stmt
 	var tail *C.char
-	rv := C.sqlite3_prepare_v2(c.db, pquery, -1, &s, &tail)
+	rv := C.sqlite3_prepare_v2_internal(c.db, pquery, -1, &s, &tail)
 	if rv != C.SQLITE_OK {
 		return nil, c.lastError()
 	}
@@ -1871,7 +1908,7 @@ func (s *SQLiteStmt) exec(ctx context.Context, args []namedValue) (driver.Result
 	}
 
 	var rowid, changes C.longlong
-	rv := C._sqlite3_step(s.s, &rowid, &changes)
+	rv := C._sqlite3_step_internal(s.s, &rowid, &changes)
 	if rv != C.SQLITE_ROW && rv != C.SQLITE_OK && rv != C.SQLITE_DONE {
 		err := s.c.lastError()
 		C.sqlite3_reset(s.s)
@@ -1943,7 +1980,7 @@ func (rc *SQLiteRows) Next(dest []driver.Value) error {
 	if rc.s.closed {
 		return io.EOF
 	}
-	rv := C.sqlite3_step(rc.s.s)
+	rv := C.sqlite3_step_internal(rc.s.s)
 	if rv == C.SQLITE_DONE {
 		return io.EOF
 	}
diff --git a/sqlite3_opt_unlock_notify.c b/sqlite3_opt_unlock_notify.c
@@ -0,0 +1,90 @@
+// Copyright (C) 2018 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+#ifdef SQLITE_ENABLE_UNLOCK_NOTIFY
+#include <sqlite3-binding.h>
+
+extern int unlock_notify_wait(sqlite3 *db);
+extern void unlock_notify_callback(void *, int);
+
+void _unlock_notify_callback(void *arg, int argc)
+{
+  unlock_notify_callback(arg, argc);
+}
+
+int
+sqlite3_step_blocking(sqlite3_stmt *stmt)
+{
+  int rv;
+  sqlite3* db;
+
+  db = sqlite3_db_handle(stmt);
+  for (;;) {
+    rv = sqlite3_step(stmt);
+    if (rv != SQLITE_LOCKED) {
+      break;
+    }
+    if (sqlite3_extended_errcode(db) != SQLITE_LOCKED_SHAREDCACHE) {
+      break;
+    }
+    rv = unlock_notify_wait(db);
+    if (rv != SQLITE_OK) {
+      break;
+    }
+    sqlite3_reset(stmt);
+  }
+
+  return rv;
+}
+
+int
+_sqlite3_step_blocking(sqlite3_stmt* stmt, long long* rowid, long long* changes)
+{
+  int rv;
+  sqlite3* db;
+
+  db = sqlite3_db_handle(stmt);
+  for (;;) {
+    rv = sqlite3_step(stmt);
+    if (rv!=SQLITE_LOCKED) {
+      break;
+    }
+    if (sqlite3_extended_errcode(db) != SQLITE_LOCKED_SHAREDCACHE) {
+      break;
+    }
+    rv = unlock_notify_wait(db);
+    if (rv != SQLITE_OK) {
+      break;
+    }
+    sqlite3_reset(stmt);
+  }
+
+  *rowid = (long long) sqlite3_last_insert_rowid(db);
+  *changes = (long long) sqlite3_changes(db);
+  return rv;
+}
+
+int
+sqlite3_prepare_v2_blocking(sqlite3 *db, const char *zSql, int nBytes, sqlite3_stmt **ppStmt, const char **pzTail)
+{
+  int rv;
+
+  for (;;) {
+    rv = sqlite3_prepare_v2(db, zSql, nBytes, ppStmt, pzTail);
+    if (rv!=SQLITE_LOCKED) {
+      break;
+    }
+    if (sqlite3_extended_errcode(db) != SQLITE_LOCKED_SHAREDCACHE) {
+      break;
+    }
+    rv = unlock_notify_wait(db);
+    if (rv != SQLITE_OK) {
+      break;
+    }
+  }
+
+  return rv;
+}
+#endif
diff --git a/sqlite3_opt_unlock_notify.go b/sqlite3_opt_unlock_notify.go
@@ -0,0 +1,53 @@
+// Copyright (C) 2018 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+// +build cgo
+// +build sqlite_unlock_notify
+
+package sqlite3
+
+/*
+#cgo CFLAGS: -DSQLITE_ENABLE_UNLOCK_NOTIFY
+
+#include <sqlite3-binding.h>
+
+extern void _unlock_notify_callback(void *arg, int argc);
+*/
+import "C"
+import (
+	"sync"
+	"unsafe"
+)
+
+type unlockNotification struct {
+	notify chan struct{}
+	lock   sync.Mutex
+}
+
+//export unlock_notify_callback
+func unlock_notify_callback(pargv unsafe.Pointer, argc C.int) {
+	argv := *(*uintptr)(pargv)
+	v := (*[1 << 30]uintptr)(unsafe.Pointer(argv))
+	for i := 0; i < int(argc); i++ {
+		un := lookupHandle(v[i]).(unlockNotification)
+		un.notify <- struct{}{}
+	}
+}
+
+var notifyMutex sync.Mutex
+
+//export unlock_notify_wait
+func unlock_notify_wait(db *C.sqlite3) C.int {
+	var un unlockNotification
+	un.notify = make(chan struct{})
+	defer close(un.notify)
+
+	argv := [1]uintptr{newHandle(nil, un)}
+	if rv := C.sqlite3_unlock_notify(db, (*[0]byte)(C._unlock_notify_callback), unsafe.Pointer(&argv)); rv != C.SQLITE_OK {
+		return rv
+	}
+	<-un.notify
+	return C.SQLITE_OK
+}
diff --git a/sqlite3_opt_unlock_notify_test.go b/sqlite3_opt_unlock_notify_test.go
@@ -0,0 +1,77 @@
+// Copyright (C) 2018 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+// +build sqlite_unlock_notify
+
+package sqlite3
+
+import (
+	"database/sql"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestTableLockedError(t *testing.T) {
+	tempFilename := TempFilename(t)
+	defer os.Remove(tempFilename)
+	dsn := fmt.Sprintf("file:%s?cache=shared&mode=rwc&_busy_timeout=%d", tempFilename, 500)
+	db, err := sql.Open("sqlite3", dsn)
+	if err != nil {
+		t.Fatal("Failed to open database:", err)
+	}
+	defer db.Close()
+
+	_, err = db.Exec("CREATE TABLE foo(id INTEGER, status INTEGER)")
+	if err != nil {
+		t.Fatal("Failed to create table:", err)
+	}
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal("Failed to begin transaction:", err)
+	}
+
+	_, err = tx.Exec("INSERT INTO foo(id, status) VALUES(1, 100)")
+	if err != nil {
+		t.Fatal("Failed to insert null:", err)
+	}
+
+	_, err = tx.Exec("UPDATE foo SET status = 200 WHERE id = 1")
+	if err != nil {
+		t.Fatal("Failed to update table:", err)
+	}
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	timer := time.NewTimer(500 * time.Millisecond)
+	go func() {
+		<-timer.C
+		err := tx.Commit()
+		if err != nil {
+			t.Fatal("Failed to commit transaction:", err)
+		}
+		wg.Done()
+	}()
+
+	rows, err := db.Query("SELECT count(*) from foo")
+	if err != nil {
+		t.Fatal("Unable to query foo table:", err)
+	}
+
+	if rows.Next() {
+		var count int
+		if err := rows.Scan(&count); err != nil {
+			t.Fatal("Failed to Scan rows", err)
+		}
+	}
+	if err := rows.Err(); err != nil {
+		t.Fatal("Failed at the call to Next:", err)
+	}
+	wg.Wait()
+
+}
