从0到1实现grpc 拦截器

前言

拦截器简单来说就是我们可以在执行业务代码前执行一段代码，和gin 的中间件的功能是一样的。我们常常在拦截器会实现限流，熔断器，auth 认证等功能。首先我们自己可以实现一个拦截器。

type handle func(ctx context.Context)
​
  var h = func(ctx context.Context) {
    fmt.Println("业务")
  }
  type interceptor func(ctx context.Context, handle2 handle)
  inter1 := func(ctx context.Context, h handle) {
    fmt.Println("1")
    h(ctx)
  }
  inter2 := func(ctx context.Context, h handle) {
    fmt.Println("2")
    h(ctx)
  }
  for _, tx := range []interceptor{inter1, inter2} {
    tx(context.Background(), h)
  }

下面是运行的结果

1
业务
2
业务

这个肯定不对呀，inter1和 inter2 是两个拦截器 h 是业务代码，业务代码应该只执行一次呀。grpc 是怎么实现的呢。

grpc 拦截器的实现

下面是我粘贴的grpc 源码

// chainUnaryClientInterceptors chains all unary client interceptors into one.
func chainUnaryClientInterceptors(cc *ClientConn) {
  interceptors := cc.dopts.chainUnaryInts
  // Prepend dopts.unaryInt to the chaining interceptors if it exists, since unaryInt will
  // be executed before any other chained interceptors.
  if cc.dopts.unaryInt != nil {
    interceptors = append([]UnaryClientInterceptor{cc.dopts.unaryInt}, interceptors...)
  }
  var chainedInt UnaryClientInterceptor
  if len(interceptors) == 0 {
    chainedInt = nil
  } else if len(interceptors) == 1 {
    chainedInt = interceptors[0]
  } else {
    chainedInt = func(ctx context.Context, method string, req, reply any, cc *ClientConn, invoker UnaryInvoker, opts ...CallOption) error {
      return interceptors[0](ctx, method, req, reply, cc, getChainUnaryInvoker(interceptors, 0, invoker), opts...)
    }
  }
  cc.dopts.unaryInt = chainedInt
}
​
// getChainUnaryInvoker recursively generate the chained unary invoker.
func getChainUnaryInvoker(interceptors []UnaryClientInterceptor, curr int, finalInvoker UnaryInvoker) UnaryInvoker {
  if curr == len(interceptors)-1 {
    return finalInvoker
  }
  return func(ctx context.Context, method string, req, reply any, cc *ClientConn, opts ...CallOption) error {
    return interceptors[curr+1](ctx, method, req, reply, cc, getChainUnaryInvoker(interceptors, curr+1, finalInvoker), opts...)
  }
}

chainUnaryClientInterceptors 这个函数主要功能就是把拦截器放在切片里，getChainUnaryInvoker 递归执行拦截器，到最后一个拦截器执行完时也就是 curr == len(interceptors)-1 执行业务函数，噢，我明白了，只要我这么做也是可以的，自己实现的拦截器。

从0到1实现拦截器

type handle func(ctx context.Context)
​
type invoke func(ctx context.Context, h handle, inter []interceptor) error
​
type interceptor func(ctx context.Context, invoker invoke, h handle) error
​
func invokers(ctx context.Context, inter []interceptor, cur int, ivk invoke) invoke {
  if cur == len(inter)-1 {
    return ivk
  }
  return func(ct context.Context, h handle, intr []interceptor) error {
    return inter[cur+1](ctx, invokers(ctx, inter, cur+1, ivk), h)
  }
}
​
func intersStart(ctx context.Context, interceptor2 []interceptor, invoker invoke) interceptor {
  if len(interceptor2) == 0 {
    return nil
  }
  if len(interceptor2) == 1 {
    return interceptor2[0]
  }
  return func(ctx context.Context, invoker invoke, h handle) error {
    return interceptor2[0](ctx, invokers(ctx, interceptor2, 0, invoker), h)
  }
}
​
func main() {
  var inters []interceptor
  inter1 := func(ctx context.Context, invoker invoke, h1 handle) error {
    fmt.Println("拦截器一")
    return invoker(ctx, h1, inters)
  }
  inter2 := func(ctx context.Context, invoker invoke, h1 handle) error {
    fmt.Println("拦截器二")
    return invoker(ctx, h1, inters)
​
  }
  inters = []interceptor{inter1, inter2}
  var invorkd = func(ctx context.Context, h handle, inter []interceptor) error {
    fmt.Println("业务")
    return nil
  }
  ctx := context.Background()
  start := intersStart(ctx, inters, invorkd)
  start(ctx, invorkd, nil)
}

执行结果

拦截器一
拦截器二
业务

我们可以看到多个拦截器，业务代码只执行了一次，也不是很难吧，最主要是 intersStart 分情况执行拦截器，多个的话 invokers 递归执行，到了最后一个拦截器执行成功了，在执行业务函数。比较绕，可以复制代码，打断点执行。