func TestTokenBucketLimiter(t *testing.T) {
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
for _, n := range []int{1, 2, 100} {
tb := jujuratelimit.NewBucketWithRate(float64(n), int64(n))
testLimiter(t, ratelimit.NewTokenBucketLimiter(tb)(e), n)
}
}
func main() {
logger := log.NewLogfmtLogger(os.Stdout)
ctx := context.Background()
c := &countService{}
var svc endpoint.Endpoint
svc = makeAddEndpoint(c)
limit := ratelimit.NewBucket(2*time.Second, 1)
svc = kitratelimit.NewTokenBucketLimiter(limit)(svc)
requestCount := expvar.NewCounter("request.count")
svc = metricsMiddleware(requestCount)(svc)
svc = loggingMiddlware(logger)(svc)
addHandler := httptransport.NewServer(
ctx,
svc,
decodeAddRequest,
encodeResponse,
httptransport.ServerBefore(beforeIDExtractor, beforePATHExtractor),
)
http.Handle("/add", addHandler)
port := os.Getenv("PORT")
logger.Log("listening on", port)
if err := http.ListenAndServe(":"+port, nil); err != nil {
logger.Log("listen.error", err)
}
}
func factory(ctx context.Context, qps int) loadbalancer.Factory {
return func(instance string) (endpoint.Endpoint, io.Closer, error) {
var e endpoint.Endpoint
e = makeUppercaseProxy(ctx, instance)
e = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{}))(e)
e = kitratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(float64(qps), int64(qps)))(e)
return e, nil, nil
}
}
func TestTokenBucketLimiter(t *testing.T) {
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
for _, n := range []int{1, 2, 100} {
testLimiter(t, ratelimit.NewTokenBucketLimiter(
ratelimit.TokenBucketLimiterRate(float64(n)),
ratelimit.TokenBucketLimiterCapacity(int64(n)),
)(e), int(n))
}
}
// New returns an AddService backed by an HTTP server living at the remote
// instance. We expect instance to come from a service discovery system, so
// likely of the form "host:port".
func New(instance string, tracer stdopentracing.Tracer, logger log.Logger) (addsvc.Service, error) {
if !strings.HasPrefix(instance, "http") {
instance = "http://" + instance
}
u, err := url.Parse(instance)
if err != nil {
return nil, err
}
// We construct a single ratelimiter middleware, to limit the total outgoing
// QPS from this client to all methods on the remote instance. We also
// construct per-endpoint circuitbreaker middlewares to demonstrate how
// that's done, although they could easily be combined into a single breaker
// for the entire remote instance, too.
limiter := ratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(100, 100))
var sumEndpoint endpoint.Endpoint
{
sumEndpoint = httptransport.NewClient(
"POST",
copyURL(u, "/sum"),
addsvc.EncodeHTTPGenericRequest,
addsvc.DecodeHTTPSumResponse,
httptransport.SetClientBefore(opentracing.FromHTTPRequest(tracer, "Sum", logger)),
).Endpoint()
sumEndpoint = opentracing.TraceClient(tracer, "Sum")(sumEndpoint)
sumEndpoint = limiter(sumEndpoint)
sumEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Sum",
Timeout: 30 * time.Second,
}))(sumEndpoint)
}
var concatEndpoint endpoint.Endpoint
{
concatEndpoint = httptransport.NewClient(
"POST",
copyURL(u, "/concat"),
addsvc.EncodeHTTPGenericRequest,
addsvc.DecodeHTTPConcatResponse,
httptransport.SetClientBefore(opentracing.FromHTTPRequest(tracer, "Concat", logger)),
).Endpoint()
concatEndpoint = opentracing.TraceClient(tracer, "Concat")(concatEndpoint)
concatEndpoint = limiter(concatEndpoint)
sumEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Concat",
Timeout: 30 * time.Second,
}))(sumEndpoint)
}
return addsvc.Endpoints{
SumEndpoint: sumEndpoint,
ConcatEndpoint: concatEndpoint,
}, nil
}
// New returns an AddService backed by a gRPC client connection. It is the
// responsibility of the caller to dial, and later close, the connection.
func New(conn *grpc.ClientConn, tracer stdopentracing.Tracer, logger log.Logger) addsvc.Service {
// We construct a single ratelimiter middleware, to limit the total outgoing
// QPS from this client to all methods on the remote instance. We also
// construct per-endpoint circuitbreaker middlewares to demonstrate how
// that's done, although they could easily be combined into a single breaker
// for the entire remote instance, too.
limiter := ratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(100, 100))
var sumEndpoint endpoint.Endpoint
{
sumEndpoint = grpctransport.NewClient(
conn,
"Add",
"Sum",
addsvc.EncodeGRPCSumRequest,
addsvc.DecodeGRPCSumResponse,
pb.SumReply{},
grpctransport.ClientBefore(opentracing.FromGRPCRequest(tracer, "Sum", logger)),
).Endpoint()
sumEndpoint = opentracing.TraceClient(tracer, "Sum")(sumEndpoint)
sumEndpoint = limiter(sumEndpoint)
sumEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Sum",
Timeout: 30 * time.Second,
}))(sumEndpoint)
}
var concatEndpoint endpoint.Endpoint
{
concatEndpoint = grpctransport.NewClient(
conn,
"Add",
"Concat",
addsvc.EncodeGRPCConcatRequest,
addsvc.DecodeGRPCConcatResponse,
pb.ConcatReply{},
grpctransport.ClientBefore(opentracing.FromGRPCRequest(tracer, "Concat", logger)),
).Endpoint()
concatEndpoint = opentracing.TraceClient(tracer, "Concat")(concatEndpoint)
concatEndpoint = limiter(concatEndpoint)
concatEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Concat",
Timeout: 30 * time.Second,
}))(concatEndpoint)
}
return addsvc.Endpoints{
SumEndpoint: sumEndpoint,
ConcatEndpoint: concatEndpoint,
}
}
func proxyingMiddleware(instances string, ctx context.Context, logger log.Logger) ServiceMiddleware {
// If instances is empty, don't proxy.
if instances == "" {
logger.Log("proxy_to", "none")
return func(next StringService) StringService { return next }
}
// Set some parameters for our client.
var (
qps = 100 // beyond which we will return an error
maxAttempts = 3 // per request, before giving up
maxTime = 250 * time.Millisecond // wallclock time, before giving up
)
// Otherwise, construct an endpoint for each instance in the list, and add
// it to a fixed set of endpoints. In a real service, rather than doing this
// by hand, you'd probably use package sd's support for your service
// discovery system.
var (
instanceList = split(instances)
subscriber sd.FixedSubscriber
)
logger.Log("proxy_to", fmt.Sprint(instanceList))
for _, instance := range instanceList {
var e endpoint.Endpoint
e = makeUppercaseProxy(ctx, instance)
e = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{}))(e)
e = ratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(float64(qps), int64(qps)))(e)
subscriber = append(subscriber, e)
}
// Now, build a single, retrying, load-balancing endpoint out of all of
// those individual endpoints.
balancer := lb.NewRoundRobin(subscriber)
retry := lb.Retry(maxAttempts, maxTime, balancer)
// And finally, return the ServiceMiddleware, implemented by proxymw.
return func(next StringService) StringService {
return proxymw{ctx, next, retry}
}
}
// New returns an AddService backed by a Thrift server described by the provided
// client. The caller is responsible for constructing the client, and eventually
// closing the underlying transport.
func New(client *thriftadd.AddServiceClient) addsvc.Service {
// We construct a single ratelimiter middleware, to limit the total outgoing
// QPS from this client to all methods on the remote instance. We also
// construct per-endpoint circuitbreaker middlewares to demonstrate how
// that's done, although they could easily be combined into a single breaker
// for the entire remote instance, too.
limiter := ratelimit.NewTokenBucketLimiter(jujuratelimit.NewBucketWithRate(100, 100))
// Thrift does not currently have tracer bindings, so we skip tracing.
var sumEndpoint endpoint.Endpoint
{
sumEndpoint = addsvc.MakeThriftSumEndpoint(client)
sumEndpoint = limiter(sumEndpoint)
sumEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Sum",
Timeout: 30 * time.Second,
}))(sumEndpoint)
}
var concatEndpoint endpoint.Endpoint
{
concatEndpoint = addsvc.MakeThriftConcatEndpoint(client)
concatEndpoint = limiter(concatEndpoint)
sumEndpoint = circuitbreaker.Gobreaker(gobreaker.NewCircuitBreaker(gobreaker.Settings{
Name: "Concat",
Timeout: 30 * time.Second,
}))(sumEndpoint)
}
return addsvc.Endpoints{
SumEndpoint: addsvc.MakeThriftSumEndpoint(client),
ConcatEndpoint: addsvc.MakeThriftConcatEndpoint(client),
}
}