331f97dbf4
* Add Twirp client/server options to preserve gRPC error details. * Pass Twirp client timeouts to the server.
54 lines
1.4 KiB
Go
54 lines
1.4 KiB
Go
package xtwirp
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/twitchtv/twirp"
|
|
)
|
|
|
|
const timeoutHeader = "X-Twirp-Timeout-Ms"
|
|
|
|
// ClientPassTimout adds context timeout as a Twirp request header.
|
|
func ClientPassTimout() twirp.ClientOption {
|
|
return twirp.WithClientInterceptors(func(fnc twirp.Method) twirp.Method {
|
|
return func(ctx context.Context, req any) (any, error) {
|
|
if deadline, ok := ctx.Deadline(); ok {
|
|
dt := time.Until(deadline)
|
|
if dt > 0 {
|
|
h, ok := twirp.HTTPRequestHeaders(ctx)
|
|
if !ok {
|
|
h = make(http.Header)
|
|
}
|
|
h.Add(timeoutHeader, strconv.FormatInt(dt.Milliseconds(), 10))
|
|
var err error
|
|
ctx, err = twirp.WithHTTPRequestHeaders(ctx, h)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
}
|
|
return fnc(ctx, req)
|
|
}
|
|
})
|
|
}
|
|
|
|
// ServerPassTimeout uses context timeout from Twirp request header.
|
|
// It requires that Twirp server handler is wrapped with PassHeadersHandler.
|
|
func ServerPassTimeout() twirp.ServerOption {
|
|
return twirp.WithServerInterceptors(func(fnc twirp.Method) twirp.Method {
|
|
return func(ctx context.Context, req any) (any, error) {
|
|
if h := GetHeaders(ctx); h != nil {
|
|
if v, err := strconv.ParseInt(h.Get(timeoutHeader), 10, 64); err == nil {
|
|
var cancel context.CancelFunc
|
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(v)*time.Millisecond)
|
|
defer cancel()
|
|
}
|
|
}
|
|
return fnc(ctx, req)
|
|
}
|
|
})
|
|
}
|