timeoutinterceptor.go 1.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. package serverinterceptors
  2. import (
  3. "context"
  4. "sync"
  5. "time"
  6. "google.golang.org/grpc"
  7. )
  8. // UnaryTimeoutInterceptor returns a func that sets timeout to incoming unary requests.
  9. func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor {
  10. return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
  11. handler grpc.UnaryHandler) (interface{}, error) {
  12. ctx, cancel := context.WithTimeout(ctx, timeout)
  13. defer cancel()
  14. var resp interface{}
  15. var err error
  16. var lock sync.Mutex
  17. done := make(chan struct{})
  18. // create channel with buffer size 1 to avoid goroutine leak
  19. panicChan := make(chan interface{}, 1)
  20. go func() {
  21. defer func() {
  22. if p := recover(); p != nil {
  23. panicChan <- p
  24. }
  25. }()
  26. lock.Lock()
  27. defer lock.Unlock()
  28. resp, err = handler(ctx, req)
  29. close(done)
  30. }()
  31. select {
  32. case p := <-panicChan:
  33. panic(p)
  34. case <-done:
  35. lock.Lock()
  36. defer lock.Unlock()
  37. return resp, err
  38. case <-ctx.Done():
  39. return nil, ctx.Err()
  40. }
  41. }
  42. }