Go的一个重要特性是可以在数据库查询时通过上下文实例取消查询(只要数据库驱动程序支持取消)。从表面上看，使用这个功能非常简单。但一旦你开始挖掘细节，就会发现有很多细微差别和相当多的陷阱……尤其是当你在web应用程序或API环境中使用这个功能时。
所以在这篇文章中，我想解释如何在web应用程序中取消数据库查询，要注意哪些奇怪的用法和边界情况，并对你可能遇到的一些情况提供解决方案。
首先，为什么要取消数据库查询？我想到了两个原因:
- 当查询的完成时间比预期的长很多时。如果出现这种情况，则说明有问题—可能是针对特定查询，也可能是针对数据库或应用程序。在这种情况下，您可能想经过一段时间后取消查询(这样资源可以得到释放，数据库连接返回到sql.DB连接池)，打印错误并返回一个500内部服务器错误给客户端。
- 当客户端在查询完成之前意外断开。出现这种情况的原因有很多，比如用户关闭浏览器选项卡或终止进程。在这个场景中，没有什么真正的“错误”，但是无需给客户端返回响应，所以您可以取消查询并释放资源。
模拟长时间运行的查询
让我们从第一种情况开始。为了演示这一点，我将编写一个非常简单的web应用程序，该程序使用pq驱动库对PostgreSQL数据库执行SELECT pg_sleep(10) SQL查询。pg_sleep(10)函数将使查询在返回之前休眠10秒，本质上模拟一个运行缓慢的查询。
package main
import (
"database/sql"
"fmt"
"log"
"net/http"
_ "github.com/lib/pq"
)
var db *sql.DB
func slowQuery() error {
_, err := db.Exec("SELECT pg_sleep(10)")
return err
}
func main() {
var err error
db, err = sql.Open("postgres", "postgres://user:pa$$word@localhost/example_db")
if err != nil {
log.Fatal(err)
}
if err = db.Ping(); err != nil {
log.Fatal(err)
}
mux := http.NewServeMux()
mux.HandleFunc("/", exampleHandler)
log.Println("Listening...")
err = http.ListenAndServe(":5000", mux)
if err != nil {
log.Fatal(err)
}
}
func exampleHandler(w http.ResponseWriter, r *http.Request) {
err := slowQuery()
if err != nil {
serverError(w, err)
return
}
fmt.Fprintln(w, "OK")
}
func serverError(w http.ResponseWriter, err error) {
log.Printf("ERROR: %s", err.Error())
http.Error(w, "Sorry, something went wrong", http.StatusInternalServerError)
}
如果您运行这段代码，然后向应用程序发出一个GET /请求，应该会发现请求挂起了10秒，然后才最终得到一个“OK”响应。像这样:
$ curl -i localhost:5000/
HTTP/1.1 200 OK
Date: Fri, 17 Apr 2020 07:46:40 GMT
Content-Length: 3
Content-Type: text/plain; charset=utf-8
OK
注意：上面的应用程序代码结构被有意地简化了。在实际的项目中，我建议使用依赖注入使得sql.DB连接池和日志对象在处理程序中可用，而不是使用全局变量。关于数据库的连接管理可以阅读Go：访问数据库代码组织方式
添加上下文超时
现在我们已经有了模拟一个长时间运行的查询代码，让我们对查询执行一个超时，这样如果查询没有在5秒内完成，它就会自动取消。
要做到这一点，我们需要:
1、使用context.withtimeout()函数来创建一个context.Context。具有5秒超时时间的上下文实例。
2、使用ExecContext()方法执行SQL查询，并将context.Context作为函数参数。
如下所示：
package main
import (
"context"
"database/sql"
"fmt"
"log"
"net/http"
"time"
_ "github.com/lib/pq"
)
var db *sql.DB
func slowQuery(ctx context.Context) error {
// 创建一个5秒超时的新子上下文，使用提供的ctx参数作为父上下文。
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
//将新的上下文作为db.ExecContext第一个参数
_, err := db.ExecContext(ctx, "SELECT pg_sleep(10)")
return err
}
...
func exampleHandler(w http.ResponseWriter, r *http.Request) {
// 将请求上下文传递给slowQuery()，以便它可以用作父上下文。
err := slowQuery(r.Context())
if err != nil {
serverError(w, err)
return
}
fmt.Fprintln(w, "OK")
}
...
这里我想强调和解释几点:
- 我们将r.Context传给slowQuery函数作为父上下文。正如后面将看到的，这个很重要，因为这意味着请求上下文上的任何取消信号都能够“扩散”到我们在ExecContext()中使用的上下文。
- defer cancel()这行很重要，是因为它确保和子上下文相关的资源在slowQuery()函数返回前被释放。如果不调用cancel()函数的话，可能会造成内存泄漏：资源不会被释放，直到父r.Context()被取消或5秒超时被命中(无论哪个最先发生)。
- 超时倒计时从使用context.withtimeout()创建子上下文的那一刻开始。如果你想对此进行更多的控制，你可以使用context.WithDeadline()函数，它允许你设置一个显式的超时时间。
我们来试试。如果再次运行应用程序并发出GET /请求，在5秒的延迟之后，你会得到这样的响应:
$ curl -i localhost:5000/
HTTP/1.1 500 Internal Server Error
Content-Type: text/plain; charset=utf-8
X-Content-Type-Options: nosniff
Date: Fri, 17 Apr 2020 08:21:14 GMT
Content-Length: 28
Sorry, something went wrong
如果你回到运行应用程序的终端窗口，应该看到类似这样的日志信息:
$ go run .
2021/09/22 10:21:07 Listening...
2021/09/22 10:21:14 ERROR: pq: canceling statement due to user request
这条日志信息可能看起来有点奇怪...直到你意识到错误消息实际上来自PostgreSQL。从这个角度看，它是有意义的：我们的web应用程序是用户在5秒后取消查询。
具体来说，在5秒后上下文超时，pq驱动发送一个取消信号给PostgreSQL数据库。然后PostgreSQL终止正在运行的查询(从而释放资源)。客户端被发送了一个500 Internal Server Error的响应，并且错误消息被记录下来，这样我们就知道什么地方出错了。
更准确地说，我们的子上下文(具有5秒超时的那个)有一个Done channel，当超时到达时，它将关闭Done channel。SQL查询运行时，我们的数据库驱动程序pq也在运行一个后台goroutine，侦听这个Done channel。如果该channel被关闭，将发送取消信号给PostgreSQL。PostgreSQL终止查询，作为对原始pq goroutine的响应，发送我们上面看到的错误消息。然后将该错误消息返回给slowQuery()函数。
处理关闭的连接
我们再试一种情况。使用curl来创建一个GET /请求，然后非常快速地(在5秒内)按Ctrl+C来取消请求。
如果您再次查看应用程序的日志，会看到另一行日志，其中包含与我们之前看到的完全相同的错误消息。
$ go run .
2021/09/22 10:21:07 Listening...
2021/09/22 10:21:14 ERROR: pq: canceling statement due to user request
2021/09/22 10:41:18 ERROR: pq: canceling statement due to user request
这里发生了什么?
在本例中，r.context请求上下文(我们在上面的代码中将其用作父上下文)被取消，因为客户端关闭了连接。
对于传入服务器的请求，如果客户端连接关闭或者ServeHTTP方法返回，请求上下文都会被canceled。
这个取消信号发送到我们的子上下文，它的Done channel关闭，并且pq驱动程序以与之前完全相同的方式终止正在运行的查询。
考虑到这一点，我们看到同样的错误消息就不足为奇了…从PostgreSQL的角度来看，发生的事情与超时时完全相同。
但从我们的web应用程序的角度来看，情况是非常不同的。客户端连接被关闭可能有许多不同的原因。从应用程序的角度来看，这并不是一个真正的错误，尽管将其作为告警记录下来可能是明智的。
幸运的是，可以通过在子上下文上调用ctx.Err()方法来区分这两个场景。如果上下文被取消(由于客户端关闭连接)，那么ctx.Err()将返回context. canceled。如果超时，它将返回context.DeadlineExceeded。如果同时达到了最后期限并且取消了上下文，那么ctx.Err()将以最先出现的错误为准。
这里需要指出的另一件重要的事情是：在PostgreSQL查询开始之前就可能发生超时/取消。例如，您可能在sql.DB连接池上设置了MaxOpenConns()最大连接数，如果达到了连接数限制，并且所有连接都在使用中，那么查询将由sql.DB“排队”，直到连接可用。在这种情况下(或任何其他导致延迟的情况下)，超时/取消很有可能在空闲数据库连接可用之前发生。在这种情况下，ExecContext()将直接返回ctx.Err()值作为响应。
如果在QueryContext()方法中，当使用Scan()处理数据时，也可能发生超时/取消。如果发生这种情况，Scan()将直接返回ctx.Err()值作为错误。据我所知，database/sql文档中并没有提到这种行为，但我在Go 1.14中碰到过这种情况。
把所有这些情况放在一起，一个明智的方法是检查“pq: canceling statement due to user request"错误，然后在slowQuery()函数返回之前用ctx.Err()中的错误包装它。
在我们的处理程序中，我们可以使用errors.Is()函数来检查slowQuery()的错误是否等于context.Canceled，并对错误进行管理。像这样:
package main
import (
"context"
"database/sql"
"errors" // New import
"fmt"
"log"
"net/http"
"time"
_ "github.com/lib/pq"
)
var db *sql.DB
func slowQuery(ctx context.Context) error {
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
_, err := db.ExecContext(ctx, "SELECT pg_sleep(10)")
//如果得到"pq: canceling statement..." 对其进行封装并返回
if err != nil && err.Error() == "pq: canceling statement due to user request" {
return fmt.Errorf("%w: %v", ctx.Err(), err)
}
return err
}
...
func exampleHandler(w http.ResponseWriter, r *http.Request) {
err := slowQuery(r.Context())
if err != nil {
//检查返回的错误是否等于context.Canceled，如果相等，记录告警。
switch {
case errors.Is(err, context.Canceled):
serverWarning(err)
default:
serverError(w, err)
}
return
}
fmt.Fprintln(w, "OK")
}
func serverWarning(err error) {
log.Printf("WARNING: %s", err.Error())
}
...
如果您现在再次运行这个应用程序，并发出两个不同的GET /请求—一个超时，另一个取消—您应该会在应用程序日志中清楚地看到不同的消息，如下所示:
$ go run .
2021/09/22 13:09:25 Listening...
2021/09/22 13:09:45 ERROR: context deadline exceeded: pq: canceling statement due to user request
2021/09/22 13:09:47 WARNING: context canceled: pq: canceling statement due to user request
其他上下文感知的方法
database/sql包为sql.DB的大部分操作提供了上下文感知变量，包括PingContext()，QueryContext()，和QueryRowContext()。我们更新main函数使用PingContext()代替Ping().
在这种情况下没有使用request context作为父context，所以需要用context.Background()创建一个空的父上下文。
...
func main() {
var err error
db, err = sql.Open("postgres", "postgres://user:pa$$word@localhost/example_db")
if err != nil {
log.Fatal(err)
}
// 用context.Background()作为父上下文创建一个10秒超时的子上下文
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// 在测试连接池时使用此选项。
if err = db.PingContext(ctx); err != nil {
log.Fatal(err)
}
mux := http.NewServeMux()
mux.HandleFunc("/", exampleHandler)
log.Println("Listening...")
err = http.ListenAndServe(":5000", mux)
if err != nil {
log.Fatal(err)
}
}
...
可以为所有请求设置一个全局超时吗?
当然，你可以在你的路由上创建并使用一些中间件，为当前的请求上下文添加一个超时，类似如下:
func setTimeout(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
defer cancel()
// 这个可以基于已有上下文创建新的请求上下文。
r = r.WithContext(ctx)
next.ServeHTTP(w, r)
})
}
如果你采用这种方法，有几件事要注意:
- 超时从创建上下文的那一刻开始，因此在数据库查询之前，处理程序中运行的任何代码都计入超时。
- 如果在一个处理程序中执行多个查询，那么它们都必须在同一时间内完成。
- 即使派生的子上下文具有不同的超时时间，该超时也将继续适用。因此，可以在子上下文中设置更早的超时，但不能使其更长。
http.TimeoutHandler
Go提供了http.TimeoutHandler中间件函数来封装你的web处理程序或router/servemux。这与上面的中间件类似，因为它在请求上下文上设置了一个超时…所以上面的警告也适用于使用这个。
然而http.TimeoutHandler还向客户端发送503 服务不可用响应和一个HTML错误消息。如果你在开发中使用这个，就不需要向客户端发送错误响应了。
context在事务中的使用
database/sql包提供了一个BeginTx()方法，可以使用它来启动上下文感知的事务。重要的是理解您提供给BeginTx()的上下文应用于整个事务。在上下文超时/取消的情况下，事务中的查询将自动回滚。
为事务中的所有查询传递相同的上下文作为参数是完全没问题的，在这种情况下，它可以确保它们在任何超时/取消之前全部(作为一个整体)完成。或者你想要每个查询单独的超时，可以创建超时时间不同的子上下文。这些子上下文必须根据BeginTx()传入的上下文创建，否则可能会存在BeginTx()中的上下文已经超时/取消并自动回滚，而代码还在执行查询。如果这种情况发生会收到“sql: transaction has already been committed or rolled back”错误。