Fix panic when closing conn during cancellable query

fixes #29
query-exec-mode
Jack Christensen 2020-04-07 19:38:21 -05:00
parent e4f3224f4c
commit 5d2be99c25
3 changed files with 31 additions and 0 deletions

View File

@ -59,6 +59,17 @@ func TestContextWatcherMultipleWatchPanics(t *testing.T) {
require.Panics(t, func() { cw.Watch(ctx2) }, "Expected panic when Watch called multiple times")
}
func TestContextWatcherUnwatchIsAlwaysSafe(t *testing.T) {
cw := ctxwatch.NewContextWatcher(func() {}, func() {})
cw.Unwatch() // unwatch when not / never watching
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cw.Watch(ctx)
cw.Unwatch()
cw.Unwatch() // double unwatch
}
func TestContextWatcherStress(t *testing.T) {
var cancelFuncCalls int64
var cleanupFuncCalls int64

View File

@ -494,6 +494,13 @@ func (pgConn *PgConn) Close(ctx context.Context) error {
defer pgConn.conn.Close()
if ctx != context.Background() {
// Close may be called while a cancellable query is in progress. This will most often be triggered by panic when
// a defer closes the connection (possibly indirectly via a transaction or a connection pool). Unwatch to end any
// previous watch. It is safe to Unwatch regardless of whether a watch is already is progress.
//
// See https://github.com/jackc/pgconn/issues/29
pgConn.contextWatcher.Unwatch()
pgConn.contextWatcher.Watch(ctx)
defer pgConn.contextWatcher.Unwatch()
}

View File

@ -1708,6 +1708,19 @@ func TestHijackAndConstruct(t *testing.T) {
ensureConnValid(t, newConn)
}
func TestConnCloseWhileCancellableQueryInProgress(t *testing.T) {
t.Parallel()
pgConn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_CONN_STRING"))
require.NoError(t, err)
ctx, _ := context.WithCancel(context.Background())
pgConn.Exec(ctx, "select n from generate_series(1,10) n")
closeCtx, _ := context.WithCancel(context.Background())
pgConn.Close(closeCtx)
}
func Example() {
pgConn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_CONN_STRING"))
if err != nil {