Goで始めるトークンバケット：基礎理論とgolang/x/time/rateの実装を読み解く

func tokenBucket_Allow() {
  l := rate.NewLimiter(1.0, 5)
  start := time.Now()
  fmt.Println("リミッター設定: 毎秒1トークン補充、バースト上限5トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  for x := range 10 {
    allowed := l.Allow()

    if allowed {
      fmt.Printf("リクエスト %d: 許可 (残りトークン: %.2f)\n", x+1, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: 拒否 (残りトークン: %.2f)\n", x+1, l.Tokens())
    }

    time.Sleep(200 * time.Millisecond)
    fmt.Printf("  待機後のトークン数: %.2f\n", l.Tokens())
  }

  fmt.Printf("経過時間: %s\n", time.Since(start).Round(time.Millisecond))
}

// 複数トークンを一度に要求するAllowNの例
func tokenBucket_AllowN() {
  l := rate.NewLimiter(5.0, 10) // 毎秒5トークン補充、バースト上限10トークン
  start := time.Now()
  fmt.Println("リミッター設定: 毎秒5トークン補充、バースト上限10トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  // 異なるトークン数でリクエスト
  tokensNeeded := []int{2, 3, 4, 5, 1, 2, 3}

  for i, n := range tokensNeeded {
    allowed := l.AllowN(time.Now(), n)

    if allowed {
      fmt.Printf("リクエスト %d: %dトークン要求 - 許可 (残りトークン: %.2f)\n", i+1, n, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: %dトークン要求 - 拒否 (残りトークン: %.2f)\n", i+1, n, l.Tokens())
    }

    time.Sleep(300 * time.Millisecond)
    fmt.Printf("  待機後のトークン数: %.2f\n", l.Tokens())
  }

  fmt.Printf("経過時間: %s\n", time.Since(start).Round(time.Millisecond))
}

func tokenBucket_Reserve() {
  l := rate.NewLimiter(2.0, 5)
  fmt.Println("リミッター設定: 毎秒2トークン補充、バースト上限5トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  for x := range 10 {
    r := l.Reserve()
    delay := r.Delay()

    if delay == 0 {
      fmt.Printf("リクエスト %d: 許可 (遅延: %v, 残りトークン: %.2f)\n", x, delay, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: 遅延発生 (遅延: %v, 残りトークン: %.2f)\n", x, delay, l.Tokens())
      // 実際には以下のようにしてdelayを待つことも可能
      // time.Sleep(delay)
    }
  }
}

func tokenBucket_ReserveN() {
  l := rate.NewLimiter(5.0, 10) // 毎秒5トークン補充、バースト上限10トークン
  fmt.Println("リミッター設定: 毎秒5トークン補充、バースト上限10トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  // 異なるトークン数でリクエスト
  tokensNeeded := []int{2, 3, 4, 5, 1, 2, 3}

  for i, n := range tokensNeeded {
    r := l.ReserveN(time.Now(), n)
    delay := r.Delay()

    if !r.OK() {
      fmt.Printf("リクエスト %d: %dトークン要求 - 予約不可 (残りトークン: %.2f)\n", i, n, l.Tokens())
    } else if delay == 0 {
      fmt.Printf("リクエスト %d: %dトークン要求 - 即時実行可能 (残りトークン: %.2f)\n", i, n, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: %dトークン要求 - %v後に実行可能 (残りトークン: %.2f)\n", i, n, delay, l.Tokens())
      // 実際には以下のようにしてdelayを待つことも可能
      // time.Sleep(delay)
    }

  }
}

func tokenBucket_Wait() {
  l := rate.NewLimiter(2.0, 5)
  fmt.Println("リミッター設定: 毎秒2トークン補充、バースト上限5トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  ctx := context.Background()

  for x := range 10 {
    err := l.Wait(ctx)

    if err == nil {
      fmt.Printf("リクエスト %d: 許可 (残りトークン: %.2f)\n", x, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: エラー %v (残りトークン: %.2f)\n", x, err, l.Tokens())
    }
  }
}

// 複数トークンを一度に要求するWaitNの例
func tokenBucket_WaitN() {
  l := rate.NewLimiter(5.0, 10) // 毎秒5トークン補充、バースト上限10トークン
  fmt.Println("リミッター設定: 毎秒5トークン補充、バースト上限10トークン")
  fmt.Println("開始時点でのトークン数:", l.Tokens())

  // 異なるトークン数でリクエスト
  tokensNeeded := []int{2, 3, 4, 5, 1, 2, 3}

  for i, n := range tokensNeeded {
    fmt.Printf("リクエスト %d: %dトークン要求 (要求前のトークン: %.2f)\n", i, n, l.Tokens())

    // 5秒のタイムアウトでコンテキストを作成
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

    // リクエスト開始時刻を記録
    start := time.Now()

    // WaitNを呼び出し、必要に応じて待機
    err := l.WaitN(ctx, n)

    // 経過時間を計算
    elapsed := time.Since(start)

    if err == nil {
      fmt.Printf("リクエスト %d: %dトークン要求 - 許可 (待機時間: %v, 残りトークン: %.2f)\n", i, n, elapsed, l.Tokens())
    } else {
      fmt.Printf("リクエスト %d: %dトークン要求 - エラー: %v (待機時間: %v, 残りトークン: %.2f)\n", i, n, err, elapsed, l.Tokens())
    }

    cancel() // コンテキストをキャンセル
  }
}

type Limiter struct {
  mu        sync.Mutex
  limit     Limit
  burst     int
  tokens    float64
  last      time.Time
  lastEvent time.Time
}

func (lim *Limiter) advance(t time.Time) (newTokens float64) {
  last := lim.last
  if t.Before(last) {
    last = t
  }

  // 経過時間に基づいてトークン数を計算
  elapsed := t.Sub(last)
  delta := lim.limit.tokensFromDuration(elapsed)
  tokens := lim.tokens + delta
  if burst := float64(lim.burst); tokens > burst {
    tokens = burst
  }
  return tokens
}

// durationFromTokens is a unit conversion function from the number of tokens to the duration
// of time it takes to accumulate them at a rate of limit tokens per second.
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
  if limit <= 0 {
    return InfDuration
  }

  duration := (tokens / float64(limit)) * float64(time.Second)

  // Cap the duration to the maximum representable int64 value, to avoid overflow.
  if duration > float64(math.MaxInt64) {
    return InfDuration
  }

  return time.Duration(duration)
}

var (
  limiter   = rate.NewLimiter(rate.Limit(1), 5)
  // ...
)

func HTTPMiddleware(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    // ...
    now := time.Now()
    allowed := limiter.AllowN(now, 1)

    // ...

    if !allowed {
      w.Header().Set("Content-Type", "application/json")
      w.WriteHeader(http.StatusTooManyRequests)
      w.Write([]byte(`{"status":"error","message":"レート制限を超えました。しばらくしてからもう一度お試しください。"}`))
      return
    }

    next.ServeHTTP(w, r)
  })
}

type Server struct{}

func (s *Server) registerEndpoints(mux *http.ServeMux) {
    // 省略
    mux.HandleFunc("GET /api/items/{id}", func(w http.ResponseWriter, r *http.Request) {
        // 省略
    }
}

func (s *Server) start() {
  mux := http.NewServeMux()
  s.registerEndpoints(mux)

  // トークンバケットレート制限の設定
  handler := HTTPMiddleware(mux)

  server := &http.Server{
    Addr:    ":8080",
    Handler: handler,
  }

  ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
  defer stop()

  s.runServer(server, ctx)
}

export const options = {
  // ...

  scenarios: {
    constant_request_rate: {
      executor: 'constant-arrival-rate',
      rate: 10,              // 1秒あたり10リクエスト
      timeUnit: '1s',        // 時間単位
      duration: '9.9s',       // 約10秒間実行
      preAllocatedVUs: 10,   // 事前に割り当てるVU（Virtual User）の数
      maxVUs: 20,            // 最大VU数
      gracefulStop: '1s',   // シナリオ終了時に進行中の処理を待つ時間
    },
  },
};


const successCount = new Counter('successCount');
const rateLimitCount = new Counter('rateLimitCount');

export default function() {
  const id = Math.floor(Math.random() * 100) + 1;
  const url = `http://localhost:8080/api/items/${id}`;

  const response = http.get(url);

  if (response.status === 200) successCount.add(1);
  if (response.status === 429) rateLimitCount.add(1);

  check(response, {
    'status is 200 or 429': (r) => r.status === 200 || r.status === 429,
  });
}

// 以下省略

type userLimiter struct {
  limiter  *rate.Limiter
  lastSeen time.Time
}

// TokenBucketLimiter はユーザーごとに独立したレート制限を管理します。
type TokenBucketLimiter struct {
  clients         map[string]*userLimiter
  mu              sync.Mutex
  rps             float64
  burst           int
  cleanupInterval time.Duration
  stop            chan struct{}
  done            chan struct{}
  closeOnce       sync.Once
}


func (t *TokenBucketLimiter) HTTPMiddleware(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    userID := strings.TrimSpace(r.Header.Get("X-User-ID"))
    if userID == "" {
      // ...
      return
    }

    now := time.Now()

    // ユーザーごとの rate limit 取得し、判定する
    limiter := t.getUserLimiter(userID, now)
    allowed := limiter.AllowN(now, 1)

    log.Printf("user=%s method=%s path=%s allowed=%t tokens=%.3f", userID, r.Method, r.URL.Path, allowed, limiter.TokensAt(now))
    if !allowed {
      writeJSON(w, http.StatusTooManyRequests, map[string]string{
        "status":  "error",
        "message": "レート制限を超えました。しばらくしてからもう一度お試しください。",
      })
      return
    }

    next.ServeHTTP(w, r)
  })
}

export const options = {
  discardResponseBodies: true,

  // ...

  scenarios: {
    per_user_rate_limit: {
      executor: 'constant-arrival-rate',
      rate: 10,
      timeUnit: '1s',
      duration: '5.9s',       // 約6秒間実行
      preAllocatedVUs: 4,
      maxVUs: 8,
      gracefulStop: '1s',
    },
  },
};

const users = ['alice', 'bob'];
// ...

export default function () {
  // 各ユーザーへ同程度のリクエストを送る。
  const iteration = exec.scenario.iterationInTest;
  const userID = users[iteration % users.length];
  const itemID = (iteration % 100) + 1;
  const response = http.get(`http://localhost:8080/api/items/${itemID}`, {
    headers: {
      'X-User-ID': userID,
    },
  });

  const isSuccess = response.status === 200;
  const isRateLimited = response.status === 429;
  if (isSuccess) userMetrics[userID].success.add(1);
  if (isRateLimited) userMetrics[userID].rateLimited.add(1);

  check(response, {
    [`${userID}: status is 200 or 429`]: (r) =>
      r.status === 200 || r.status === 429,
  });
}


// 以下省略