服务治理：熔断 / 限流 / 降级原理与 gRPC 接入

package breaker

import (
    "errors"
    "sync"
    "time"
)

// State 熔断器状态
type State int

const (
    StateClosed State = iota
    StateOpen
    StateHalfOpen
)

var (
    ErrCircuitOpen = errors.New("circuit breaker is open")
)

// Breaker 简单熔断器
type Breaker struct {
    mu              sync.Mutex
    state           State
    failureThreshold int     // 错误数阈值
    failureCount    int     // 当前窗口错误数
    requestCount    int     // 当前窗口总请求数
    lastStateChange time.Time
    openTimeout     time.Duration // Open 状态持续多久后转 HalfOpen
    halfOpenMaxCalls int          // HalfOpen 状态允许试探的请求数
    halfOpenCalls    int
}

func New(failureThreshold int, openTimeout time.Duration) *Breaker {
    return &Breaker{
        state:            StateClosed,
        failureThreshold: failureThreshold,
        openTimeout:      openTimeout,
        halfOpenMaxCalls: 5, // 默认放 5 个试探请求
    }
}

// Allow 判断请求是否允许通过
func (b *Breaker) Allow() error {
    b.mu.Lock()
    defer b.mu.Unlock()

    switch b.state {
    case StateClosed:
        // 步骤 1：闭合态直接放行
        return nil
    case StateOpen:
        // 步骤 2：开放态，只有超时后才转半开并放行一个试探
        if time.Since(b.lastStateChange) > b.openTimeout {
            b.state = StateHalfOpen
            b.halfOpenCalls = 0
            return nil
        }
        return ErrCircuitOpen
    case StateHalfOpen:
        // 步骤 3：半开态限制试探请求数，超了就拒绝
        if b.halfOpenCalls < b.halfOpenMaxCalls {
            b.halfOpenCalls++
            return nil
        }
        return ErrCircuitOpen
    }
    return nil
}

// MarkSuccess 标记一次成功
func (b *Breaker) MarkSuccess() {
    b.mu.Lock()
    defer b.mu.Unlock()
    if b.state == StateHalfOpen {
        // 步骤 1：半开成功 → 转 Closed，清空计数
        b.state = StateClosed
        b.failureCount = 0
        b.requestCount = 0
    }
}

// MarkFailure 标记一次失败
func (b *Breaker) MarkFailure() {
    b.mu.Lock()
    defer b.mu.Unlock()
    b.failureCount++
    b.requestCount++

    if b.state == StateHalfOpen {
        // 步骤 1：半开失败 → 立刻转 Open，记录时间
        b.state = StateOpen
        b.lastStateChange = time.Now()
        return
    }

    // 步骤 2：闭合态错误数达阈值 → 转 Open
    if b.failureCount >= b.failureThreshold {
        b.state = StateOpen
        b.lastStateChange = time.Now()
    }
}

package counter

import "sync"

// CounterLimiter 计数器限流器，限制并发数
type CounterLimiter struct {
    mu       sync.Mutex
    current  int // 当前并发数
    maxConns int // 最大并发数
}

func New(maxConns int) *CounterLimiter {
    return &CounterLimiter{maxConns: maxConns}
}

// Allow 尝试获取一个并发槽位
func (l *CounterLimiter) Allow() bool {
    l.mu.Lock()
    defer l.mu.Unlock()
    // 步骤 1：已达上限直接拒绝
    if l.current >= l.maxConns {
        return false
    }
    // 步骤 2：占用一个槽位
    l.current++
    return true
}

// Release 处理完毕后释放槽位
func (l *CounterLimiter) Release() {
    l.mu.Lock()
    defer l.mu.Unlock()
    if l.current > 0 {
        l.current--
    }
}

package slidingwindow

import (
    "sync"
    "time"
)

// SlidingWindowLimiter 滑动窗口限流器
type SlidingWindowLimiter struct {
    mu        sync.Mutex
    window    time.Duration // 窗口大小，如 1s
    threshold int           // 窗口内允许的最大请求数
    requests  []time.Time   // 窗口内的请求时间戳
}

func New(window time.Duration, threshold int) *SlidingWindowLimiter {
    return &SlidingWindowLimiter{
        window:    window,
        threshold: threshold,
    }
}

// Allow 检查是否允许通过
func (l *SlidingWindowLimiter) Allow() bool {
    l.mu.Lock()
    defer l.mu.Unlock()

    now := time.Now()
    // 步骤 1：清理窗口外的旧请求
    boundary := now.Add(-l.window)
    idx := 0
    for ; idx < len(l.requests); idx++ {
        if l.requests[idx].After(boundary) {
            break
        }
    }
    l.requests = l.requests[idx:]

    // 步骤 2：超过阈值则拒绝
    if len(l.requests) >= l.threshold {
        return false
    }
    // 步骤 3：记录本次请求并放行
    l.requests = append(l.requests, now)
    return true
}

package tokenbucket

import (
    "sync"
    "time"
)

// TokenBucket 令牌桶限流器
type TokenBucket struct {
    mu         sync.Mutex
    capacity   int64         // 桶容量（最大令牌数）
    tokens     int64         // 当前令牌数
    rate       int64         // 每秒生成的令牌数
    lastRefill time.Time     // 上次补充令牌时间
}

func New(capacity, rate int64) *TokenBucket {
    return &TokenBucket{
        capacity:   capacity,
        tokens:     capacity, // 启动时桶是满的，允许初始突发
        rate:       rate,
        lastRefill: time.Now(),
    }
}

// Allow 尝试获取一个令牌
func (tb *TokenBucket) Allow() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()

    // 步骤 1：按经过时间补充令牌
    now := time.Now()
    elapsed := now.Sub(tb.lastRefill).Seconds()
    refill := int64(elapsed * float64(tb.rate))
    if refill > 0 {
        tb.tokens += refill
        if tb.tokens > tb.capacity {
            tb.tokens = tb.capacity // 不能超过桶容量
        }
        tb.lastRefill = now
    }

    // 步骤 2：有令牌就消费并放行，否则拒绝
    if tb.tokens > 0 {
        tb.tokens--
        return true
    }
    return false
}

import (
    "context"
    "google.golang.org/grpc"
)

// LogInterceptor 简单日志拦截器
func LogInterceptor(
    ctx context.Context,
    req any,
    info *grpc.UnaryServerInfo,
    handler grpc.UnaryHandler,
) (any, error) {
    // 前置逻辑：记录调用开始
    start := time.Now()
    // 调用真正的 handler
    resp, err := handler(ctx, req)
    // 后置逻辑：记录耗时和错误
    log.Printf("method=%s cost=%v err=%v", info.FullMethod, time.Since(start), err)
    return resp, err
}

// 注册时用 ChainUnaryInterceptor 串联多个拦截器
server := grpc.NewServer(
    grpc.ChainUnaryInterceptor(LogInterceptor, /* ... */),
)

// Limiter 我们之前实现的限流器接口
type Limiter interface {
    Allow(key string) bool
    Release(key string)
}

// GRPCLimitInterceptor 全应用统一限流
func GRPCLimitInterceptor(l Limiter) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
        // 步骤 1：用固定 key 对整个应用一起限流
        const key = "global"
        if !l.Allow(key) {
            // 返回 gRPC 标准限流错误码
            return nil, status.Error(codes.ResourceExhausted, "rate limited")
        }
        // 步骤 2：处理完释放槽位
        defer l.Release(key)
        return handler(ctx, req)
    }
}

// ServiceLimitInterceptor 按 FullMethod 限流，粒度更细
func ServiceLimitInterceptor(l Limiter) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
        // 步骤 1：用方法名作为 key，针对每个方法独立限流
        if !l.Allow(info.FullMethod) {
            return nil, status.Error(codes.ResourceExhausted, "rate limited")
        }
        // 步骤 2：处理完释放
        defer l.Release(info.FullMethod)
        return handler(ctx, req)
    }
}

// BizLimitInterceptor 按业务字段（如 UID）限流
// 通过对 req 类型断言取出业务字段
type uidRequest interface {
    GetUid() int64
}

func BizLimitInterceptor(l Limiter) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
        // 步骤 1：类型断言，取出 UID 作为 key
        if r, ok := req.(uidRequest); ok {
            key := strconv.FormatInt(r.GetUid(), 10)
            if !l.Allow(key) {
                return nil, status.Error(codes.ResourceExhausted, "rate limited")
            }
            // 步骤 2：处理完释放
            defer l.Release(key)
        }
        // 没实现 GetUid 的请求直接放行（会浪费一点 CPU 检查）
        return handler(ctx, req)
    }
}

// RateLimitedUserService 用装饰器模式给 UserService 加限流
type RateLimitedUserService struct {
    UserServiceServer
    limiter Limiter
}

func (r *RateLimitedUserService) GetProfile(ctx context.Context, req *GetProfileReq) (*ProfileResp, error) {
    // 步骤 1：按 UID 限流
    key := strconv.FormatInt(req.GetUid(), 10)
    if !r.limiter.Allow(key) {
        return nil, status.Error(codes.ResourceExhausted, "rate limited")
    }
    // 步骤 2：放行到被装饰的实现
    defer r.limiter.Release(key)
    return r.UserServiceServer.GetProfile(ctx, req)
}

import (
    "github.com/go-kratos/aegis/circuitbreaker"
)

// CircuitBreakerInterceptor 熔断拦截器
func CircuitBreakerInterceptor(cb *circuitbreaker.Breaker) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
        // 步骤 1：Allow 返回 error 表示熔断已开启，直接拒绝
        if err := cb.Allow(); err != nil {
            return nil, status.Error(codes.Unavailable, "circuit breaker open")
        }
        // 步骤 2：调用业务，拿到结果与错误
        resp, err := handler(ctx, req)
        // 步骤 3：关键——成对上报成功/失败，否则熔断状态永不更新
        if err != nil {
            cb.MarkFailed()
        } else {
            cb.MarkSuccess()
        }
        return resp, err
    }
}

// 业务代码里结合限流器状态做降级
func (s *UserService) GetProfile(ctx context.Context, uid int64) (*Profile, error) {
    // 步骤 1：触发限流时，缓存未命中直接返回（不走 DB）
    if s.limiter.IsTriggered() {
        p, err := s.cache.Get(ctx, uid)
        if err == nil {
            return p, nil
        }
        // 缓存未命中，降级：返回错误或默认值
        return nil, err.ErrLimited
    }
    // 步骤 2：正常路径
    return s.cacheAside.Get(ctx, uid)
}

ctx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
resp, err := client.GetProfile(ctx, req)

func LogInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
    // 步骤 1：记录开始时间
    start := time.Now()
    // 步骤 2：调用 handler
    resp, err := handler(ctx, req)
    // 步骤 3：记录耗时、错误码、方法名
    log.Printf("method=%s cost=%v err=%v", info.FullMethod, time.Since(start), err)
    return resp, err
}