123456789101112131415161718192021222324252627282930313233343536373839404142434445464748 |
- package serverinterceptors
- import (
- "context"
- "sync"
- "time"
- "google.golang.org/grpc"
- )
- // UnaryTimeoutInterceptor returns a func that sets timeout to incoming unary requests.
- func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor {
- return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
- handler grpc.UnaryHandler) (interface{}, error) {
- ctx, cancel := context.WithTimeout(ctx, timeout)
- defer cancel()
- var resp interface{}
- var err error
- var lock sync.Mutex
- done := make(chan struct{})
- // create channel with buffer size 1 to avoid goroutine leak
- panicChan := make(chan interface{}, 1)
- go func() {
- defer func() {
- if p := recover(); p != nil {
- panicChan <- p
- }
- }()
- lock.Lock()
- defer lock.Unlock()
- resp, err = handler(ctx, req)
- close(done)
- }()
- select {
- case p := <-panicChan:
- panic(p)
- case <-done:
- lock.Lock()
- defer lock.Unlock()
- return resp, err
- case <-ctx.Done():
- return nil, ctx.Err()
- }
- }
- }
|