Question

添加上下文，避免操作阻塞。

package main

import (
	"context"
	"database/sql"
	"time"
	"log"

	_ "github.com/mattn/go-sqlite3"
)


var (
	db *sql.DB
	
	ctx = context.Background()
	timeout = time.Millisecond
)

func ping() {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var name string
	err := db.QueryRowContext(ctx, `SELECT name FROM user`).Scan(&name)

	if err == context.DeadlineExceeded {
		println("timeout!")
	} else if err != nil { 
		log.Println(err) 
	} else {
		println(name)
	}

	// select {
	// case <- ctx.Done(): println("timeout")
	// default:
	// }	
}

func main() {
	db, _ = sql.Open("sqlite3", "./test.db")
	ping()
}

源码剖析

所有操作都有其上下文版本。

// database/sql/sql.go

func (db *DB) Query(query string, args ...any) (*Rows, error) {
	return db.QueryContext(context.Background(), query, args...)
}

func (db *DB) QueryContext(ctx context.Context, ...) (*Rows, error) {
	return db.query(ctx, query, args, alwaysNewConn)
}

func (db *DB) query(ctx context.Context, ...) (*Rows, error) {
	dc, err := db.conn(ctx, strategy)
	return db.queryDC(ctx, nil, dc, dc.releaseConn, query, args)
}

func (db *DB) queryDC(ctx, txctx context.Context, ...) (*Rows, error) {
	rowsi, err = ctxDriverQuery(ctx, queryerCtx, queryer, query, nvdargs)
}

类似 ctxDriverQuery 这样的包装函数，专门针对上下文做出反应。

// database/sql/ctxutil.go

func ctxDriverQuery(ctx context.Context, ...) (driver.Rows, error) {
	select {
	default:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
    
	return queryer.Query(query, dargs)
}

在操作内部，也会建立上下文应对机制。

// database/sql/sql.go

func (db *DB) queryDC(ctx, txctx context.Context, ...) (*Rows, error) {
	rows.initContextClose(ctx, txctx)
}

func (rs *Rows) initContextClose(ctx, txctx context.Context) {
	ctx, rs.cancel = context.WithCancel(ctx)
	go rs.awaitDone(ctx, txctx)
}

func (rs *Rows) awaitDone(ctx, txctx context.Context) {
	select {
	case <-ctx.Done():
	case <-txctxDone:
	}
	rs.close(ctx.Err())
}