如何在Go中使用context实现请求限制

在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实现了请求计数。

后端开发标签