From 7e1a61dbcdb71e3220e5cc54cd2ef7d32b0b9b67 Mon Sep 17 00:00:00 2001
From: Andrii Zavorotnii <andrii.zavorotnii@gmail.com>
Date: Thu, 5 Sep 2019 15:26:33 -0700
Subject: Fix context cancellation racy handling

[why]
Context cancellation goroutine is not in sync with Next() method lifetime.
It leads to sql.ErrNoRows instead of context.Canceled often (easy to reproduce).
It leads to interruption of next query executed on same connection (harder to reproduce).

[how]
Do query in goroutine, wait when interruption done.

[testing]
Add unit test that reproduces error cases.
---
 sqlite3_go18_test.go | 88 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

(limited to 'sqlite3_go18_test.go')

diff --git a/sqlite3_go18_test.go b/sqlite3_go18_test.go
index c9e79e7..37117b0 100644
--- a/sqlite3_go18_test.go
+++ b/sqlite3_go18_test.go
@@ -14,6 +14,7 @@ import (
 	"io/ioutil"
 	"math/rand"
 	"os"
+	"sync"
 	"testing"
 	"time"
 )
@@ -135,6 +136,93 @@ func TestShortTimeout(t *testing.T) {
 	}
 }
 
+func TestQueryRowContextCancel(t *testing.T) {
+	srcTempFilename := TempFilename(t)
+	defer os.Remove(srcTempFilename)
+
+	db, err := sql.Open("sqlite3", srcTempFilename)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	initDatabase(t, db, 100)
+
+	const query = `SELECT key_id FROM test_table ORDER BY key2 ASC`
+	var keyID string
+	unexpectedErrors := make(map[string]int)
+	for i := 0; i < 10000; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+		row := db.QueryRowContext(ctx, query)
+
+		cancel()
+		// it is fine to get "nil" as context cancellation can be handled with delay
+		if err := row.Scan(&keyID); err != nil && err != context.Canceled {
+			if err.Error() == "sql: Rows are closed" {
+				// see https://github.com/golang/go/issues/24431
+				// fixed in 1.11.1 to properly return context error
+				continue
+			}
+			unexpectedErrors[err.Error()]++
+		}
+	}
+	for errText, count := range unexpectedErrors {
+		t.Error(errText, count)
+	}
+}
+
+func TestQueryRowContextCancelParallel(t *testing.T) {
+	srcTempFilename := TempFilename(t)
+	defer os.Remove(srcTempFilename)
+
+	db, err := sql.Open("sqlite3", srcTempFilename)
+	if err != nil {
+		t.Fatal(err)
+	}
+	db.SetMaxOpenConns(10)
+	db.SetMaxIdleConns(5)
+
+	defer db.Close()
+	initDatabase(t, db, 100)
+
+	const query = `SELECT key_id FROM test_table ORDER BY key2 ASC`
+	wg := sync.WaitGroup{}
+	defer wg.Wait()
+
+	testCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			var keyID string
+			for {
+				select {
+				case <-testCtx.Done():
+					return
+				default:
+				}
+				ctx, cancel := context.WithCancel(context.Background())
+				row := db.QueryRowContext(ctx, query)
+
+				cancel()
+				_ = row.Scan(&keyID) // see TestQueryRowContextCancel
+			}
+		}()
+	}
+
+	var keyID string
+	for i := 0; i < 10000; i++ {
+		// note that testCtx is not cancelled during query execution
+		row := db.QueryRowContext(testCtx, query)
+
+		if err := row.Scan(&keyID); err != nil {
+			t.Fatal(i, err)
+		}
+	}
+}
+
 func TestExecCancel(t *testing.T) {
 	db, err := sql.Open("sqlite3", ":memory:")
 	if err != nil {
-- 
cgit v1.2.3