-
Notifications
You must be signed in to change notification settings - Fork 5
/
middleware.go
80 lines (68 loc) · 1.99 KB
/
middleware.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package timeoutmiddleware
import (
"context"
"github.com/arquivei/foundationkit/errors"
"github.com/go-kit/kit/endpoint"
)
// New returns a new timeout middleware.
//
// After timeout is reached, if the middleware is configured to wait,
// it will just cancel the context and wait for next endpoint to return.
// But if the middleware is configured to not wait, it will run the next endpoint
// inside a go-routine and return error as soon as the context is canceled.
func New(c Config) (endpoint.Middleware, error) {
return func(next endpoint.Endpoint) endpoint.Endpoint {
// Timeout is disabled
if c.Timeout <= 0 {
return next
}
return func(ctx context.Context, request interface{}) (response interface{}, err error) {
ctx, cancel := context.WithTimeout(ctx, c.Timeout)
defer cancel()
// Override error code and severity based on the context
defer func() {
if err != nil && ctx.Err() != nil {
err = errors.E(err, c.ErrorSeverity, c.ErrorCode)
}
}()
if c.Wait {
return next(ctx, request)
}
return nextNoWait(ctx, next, request)
}
}, nil
}
// nextNoWait runs next but don't wait for a response in case of canceled context
func nextNoWait(ctx context.Context, next endpoint.Endpoint, request interface{}) (interface{}, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case r := <-runNextAsync(ctx, next, request):
return r.response, r.err
}
}
type asyncResult struct {
response interface{}
err error
}
// runNextAsync executes next inside a go-routine and returns the result in a channel.
func runNextAsync(ctx context.Context, next endpoint.Endpoint, request interface{}) <-chan asyncResult {
c := make(chan asyncResult)
go func() {
defer close(c)
// Panics in go-routines must be captured inside the go-routine
err := errors.DontPanic(func() {
response, err := next(ctx, request)
c <- asyncResult{
response: response,
err: err,
}
})
if err != nil {
c <- asyncResult{
err: err,
}
}
}()
return c
}