在Go语言中,context是一个非常常用的包,它为我们提供了一种在不同函数间传递请求附加值的方式,它可以用于跟踪请求、超时、取消信号、截止日期等等。本篇文章将会讲解如何使用context来实现请求限制,我们将通过两种方法来实现:
1.使用context.WithDeadline
context.WithDeadline函数用于创建一个带有时限的上下文对象,这个上下文对象会在截止日期到来时自动取消,我们可以使用这个函数来限制请求的执行时间。
我们首先定义一个函数来模拟一个需要执行5秒钟才能完成的任务,然后使用context.WithDeadline来限制这个任务的执行时间。
使用context.WithDeadline限制请求时间
示例代码
package main
import (
"context"
"fmt"
"time"
)
func task(ctx context.Context) {
deadline, ok := ctx.Deadline()
if ok {
fmt.Printf("Deadline: %s\n", deadline)
}
select {
case <-ctx.Done():
fmt.Println("task canceled")
return
case <-time.After(5 * time.Second):
fmt.Println("task completed")
}
}
func main() {
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Second))
defer cancel()
task(ctx)
}
我们首先定义了一个函数task,它接收一个上下文对象作为参数,在这个函数中,我们使用select来监听ctx.Done()这个channel,如果任务被取消了,就会从这个channel中读取到值,然后我们就可以结束这个函数。如果任务没有被取消,我们就使用time.After来模拟一个需要5秒才能完成的任务,等待5秒后完成任务。
在main函数中,我们使用context.WithDeadline来创建一个带有2秒的时限的上下文对象,然后把这个上下文对象传给task函数,这样task函数就只有2秒的时间来完成任务,如果2秒内任务没有完成,这个上下文对象就会被取消,这个任务也会被中止。
示例输出
当运行这个程序时,我们会发现在2秒后,“task canceled”这个字符串被输出了,这说明我们成功地使用context.WithDeadline限制了这个任务的执行时间。
使用context.WithValue实现简单请求计数
另一种使用context的方法是使用context.WithValue函数来传递请求附加值,我们可以使用这个函数来实现简单的请求计数。
我们首先定义一个中间件函数countRequests,在这个函数中,我们使用context.Value来获取一个整数类型的请求计数器,如果这个计数器不存在,我们就创建一个新的计数器,并把它附加到上下文对象中。然后我们输出这个计数器的值,并把这个上下文对象传给下一个处理函数。
使用context.WithValue实现请求计数
示例代码
package main
import (
"context"
"fmt"
"net/http"
)
type key int
const (
requestCounterKey key = iota
)
func countRequests(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
// Get the request counter from the context
counter, ok := ctx.Value(requestCounterKey).(int)
if !ok {
counter = 0
}
// Increment the request counter
counter++
fmt.Printf("Request count: %d\n", counter)
// Put the request counter back in the context
ctx = context.WithValue(ctx, requestCounterKey, counter)
// Call the next handler in the chain
next.ServeHTTP(w, r.WithContext(ctx))
})
}
func helloWorld(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
// Get the request counter from the context
counter, ok := ctx.Value(requestCounterKey).(int)
if !ok {
counter = 0
}
// Greet the client and show them their request count
fmt.Fprintf(w, "Hello, world! This is request #%d.", counter)
}
func main() {
http.Handle("/", countRequests(http.HandlerFunc(helloWorld)))
http.ListenAndServe(":8080", nil)
}
我们首先定义了一个整数类型的常量requestCounterKey来代表存储请求计数器的键,然后我们定义了一个中间件函数countRequests,它接收一个http.Handler类型的参数,并返回一个新的http.Handler类型的值。在这个中间件函数中,我们使用上下文对象中存储的计数器来计算当前请求数量,然后把这个计数器加1,并把它放回上下文对象中,然后把这个上下文对象传给下一个处理函数。
在main函数中,我们使用http.Handle来注册一个处理器函数countRequests来处理请求,并把这个处理器函数包装在helloWorld这个处理函数中,这样每当有新的请求到达时,它们就会被countRequests处理器函数处理,然后才会交给helloWorld处理函数处理。在helloWorld处理函数中,我们使用上下文对象中存储的计数器来计算当前请求数量,并以此为依据返回一个问候信息给客户端。
示例输出
当运行这个程序时,我们可以通过访问http://localhost:8080来测试这个程序。每当有新的请求到达时,它们就会被countRequests处理器函数处理,并输出当前的请求数量。这说明我们成功地使用context实现了请求计数。