Golang中如何实现不带ctx/context的限流

Golang中如何实现不带ctx/context的限流 大家好,

我是Go语言的新手,学习它只是因为需要实现限流Webhook。我写了下面这段代码:

func rateLimiter(next func(w http.ResponseWriter, r *http.Request, _ httprouter.Params)) httprouter.Handle {
	qps := 1
	burst := 1
	// ctx := context.TODO()
	limit := rate.Limit(qps)
	limiter := rate.NewLimiter(limit, burst)	
	return httprouter.Handle(func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
		// for i := uint(1); i < 100; i++ {
			if !limiter.Allow() {
				if r.Method == "POST" {
					log.Println("Blocked by rate limit!")
				}
				// _ = limiter.Wait(ctx)	
				return
			} // else {
				// time.Sleep(time.Duration(rand.Intn(500)) * time.Millisecond)
			next(w, r, nil)
			// }
		// }
	})
}

我的主管用VSCode和Bito(AI工具)检查了这段代码。Bito建议移除rateLimiter函数中的ctx(如上所示),以及for循环等,具体请看代码中的//注释部分。

我的理解是,ctx是用来区分不同请求的令牌,移除它会导致程序只接受一个请求而阻塞其他请求。那么,如果两个请求同时发生,是不是只有一个会被接受呢?

感谢您的帮助。


更多关于Golang中如何实现不带ctx/context的限流的实战教程也可以访问 https://www.itying.com/category-94-b0.html

1 回复

更多关于Golang中如何实现不带ctx/context的限流的实战系列教程也可以访问 https://www.itying.com/category-94-b0.html


在Go语言中,使用rate.Limiter实现不带context的限流是完全可行的。你的代码基本正确,但需要理解limiter.Allow()的工作原理。

rate.Limiter是线程安全的,多个goroutine可以并发调用它的方法。limiter.Allow()会检查当前是否允许执行,如果允许则消耗一个令牌并返回true,否则返回false。它不依赖于context来区分请求。

你的代码示例:

package main

import (
    "log"
    "net/http"
    "time"

    "github.com/julienschmidt/httprouter"
    "golang.org/x/time/rate"
)

func rateLimiter(next func(w http.ResponseWriter, r *http.Request, _ httprouter.Params)) httprouter.Handle {
    qps := 1  // 每秒1个请求
    burst := 1 // 突发容量为1
    limit := rate.Limit(qps)
    limiter := rate.NewLimiter(limit, burst)
    
    return httprouter.Handle(func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
        if !limiter.Allow() {
            log.Printf("Rate limit exceeded for %s %s", r.Method, r.URL.Path)
            http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
            return
        }
        next(w, r, nil)
    })
}

func handler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
    w.Write([]byte("Request processed at " + time.Now().Format("15:04:05.000")))
}

func main() {
    router := httprouter.New()
    router.POST("/webhook", rateLimiter(handler))
    
    log.Println("Server starting on :8080")
    log.Fatal(http.ListenAndServe(":8080", router))
}

并发测试示例:

package main

import (
    "fmt"
    "sync"
    "time"
    "golang.org/x/time/rate"
)

func main() {
    limiter := rate.NewLimiter(1, 1) // 1 QPS, burst 1
    
    var wg sync.WaitGroup
    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            
            if limiter.Allow() {
                fmt.Printf("Goroutine %d: Allowed at %v\n", 
                    id, time.Now().Format("15:04:05.000"))
            } else {
                fmt.Printf("Goroutine %d: Denied at %v\n", 
                    id, time.Now().Format("15:04:05.000"))
            }
        }(i)
    }
    
    wg.Wait()
    time.Sleep(2 * time.Second)
    
    // 测试时间窗口内的请求
    fmt.Println("\nTesting time-based limiting:")
    for i := 1; i <= 3; i++ {
        if limiter.Allow() {
            fmt.Printf("Request %d: Allowed\n", i)
        } else {
            fmt.Printf("Request %d: Denied\n", i)
        }
        time.Sleep(500 * time.Millisecond)
    }
}

输出可能类似于:

Goroutine 2: Denied at 14:30:25.123
Goroutine 1: Allowed at 14:30:25.123
Goroutine 3: Denied at 14:30:25.123
Goroutine 4: Denied at 14:30:25.123
Goroutine 5: Denied at 14:30:25.123

Testing time-based limiting:
Request 1: Allowed
Request 2: Denied
Request 3: Allowed

关键点:

  1. limiter.Allow()是原子的,多个goroutine同时调用时,只有一个会成功(当burst=1时)
  2. 令牌桶算法会自动按照设定的速率补充令牌
  3. 不需要context来区分请求,limiter本身管理令牌状态
  4. 对于Webhook场景,这种实现是合适的,超过限制的请求会立即被拒绝

你的理解有误:context在这里不是用来区分请求的,而是用于超时或取消控制。limiter.Wait(ctx)会阻塞直到获取令牌或context被取消,而limiter.Allow()是立即返回的非阻塞方法。

回到顶部