github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/clientinterceptors/timeoutinterceptor.go (about)

     1  package clientinterceptors
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"google.golang.org/grpc"
     8  )
     9  
    10  // TimeoutInterceptor is an interceptor that controls timeout.
    11  func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
    12  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
    13  		invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    14  		if timeout <= 0 {
    15  			return invoker(ctx, method, req, reply, cc, opts...)
    16  		}
    17  
    18  		ctx, cancel := context.WithTimeout(ctx, timeout)
    19  		defer cancel()
    20  
    21  		return invoker(ctx, method, req, reply, cc, opts...)
    22  	}
    23  }