Now that we have built a complete service, we are going to see how easy it is to add middleware to our endpoints in order to extend the service without touching the actual implementations themselves.
In real-world services, it is sensible to limit the number of requests it will attempt to handle so that the service doesn't get overwhelmed. This can happen if the process needs more memory than is available, or we might notice performance degradation if it eats up too much of the CPU. In a micro-service architecture, the strategy to solving these problems is to add another node and spread the load, which means that we want each individual instance to be rate limited.
Since we are providing the client, we should add rate limiting there, which would prevent too many requests from getting on the network. But it is also sensible to add rate limiting to the server in case many clients are trying to access the same services at the same time. Luckily, endpoints in Go kit are used for both the client and server, so we can use the same code to add middleware in both places.
We are going to add a Token Bucket-based rate limiter, which you can read more about at https://en.wikipedia.org/wiki/Token_bucket. The guys at Juju have written a Go implementation that we can use by importing github.com/juju/ratelimit, and Go kit has middleware built for this very implementation, which will save us a lot of time and effort.
The general idea is that we have a bucket of tokens, and each request will need a token in order to do its work. If there are no tokens in the bucket, we have reached our limit and the request cannot be completed. Buckets refill over time at a specific interval.
Import github.com/juju/ratelimit and before we create our hashEndpoint, insert the following code:
rlbucket := ratelimit.NewBucket(1*time.Second, 5)
The NewBucket function creates a new rate limiting bucket that will refill at a rate of one token per second, up to a maximum of five tokens. These numbers are pretty silly for our case, but we want to be able to reach our limits manually during the development.
Since the Go kit ratelimit package has the same name as the Juju one, we are going to need to import it with a different name:
import ratelimitkit "github.com/go-kit/kit/ratelimit"
Endpoint middleware in Go kit is specified by the endpoint.Middleware function type:
type Middleware func(Endpoint) Endpoint
A piece of middleware is simply a function that takes Endpoint and returns Endpoint. Remember that Endpoint is also a function:
type Endpoint func(ctx context.Context, request
interface{}) (response interface{}, err error)
This gets a little confusing, but they are the same as the wrappers we built for http.HandlerFunc. A middleware function returns an Endpoint function that does something before and/or after calling the Endpoint being wrapped. The arguments passed into the function that returns the Middleware are closured in, which means that they are available to the inner code (via closures) without the state having to be stored anywhere else.
We are going to use the NewTokenBucketLimiter middleware from Go kit's ratelimit package, and if we take a look at the code, we'll see how it uses closures and returns functions to inject a call to the token bucket's TakeAvailable method before passing execution to the next endpoint:
func NewTokenBucketLimiter(tb *ratelimit.Bucket)
endpoint.Middleware {
return func(next endpoint.Endpoint) endpoint.Endpoint {
return func(ctx context.Context, request interface{})
(interface{}, error) {
if tb.TakeAvailable(1) == 0 {
return nil, ErrLimited
}
return next(ctx, request)
}
}
}
A pattern has emerged within Go kit where you obtain the endpoint and then put all middleware adaptations inside their own block immediately afterwards. The returned function is given the endpoint when it is called, and the same variable is overwritten with the result.
For a simple example, consider this code:
e := getEndpoint(srv)
{
e = getSomeMiddleware()(e)
e = getLoggingMiddleware(logger)(e)
e = getAnotherMiddleware(something)(e)
}
We will now do this for our endpoints; update the code inside the main function to add the rate limiting middleware:
hashEndpoint := vault.MakeHashEndpoint(srv)
{
hashEndpoint = ratelimitkit.NewTokenBucketLimiter
(rlbucket)(hashEndpoint)
}
validateEndpoint := vault.MakeValidateEndpoint(srv)
{
validateEndpoint = ratelimitkit.NewTokenBucketLimiter
(rlbucket)(validateEndpoint)
}
endpoints := vault.Endpoints{
HashEndpoint: hashEndpoint,
ValidateEndpoint: validateEndpoint,
}
There's nothing much to change here; we're just updating the hashEndpoint and validateEndpoint variables before assigning them to the vault.Endpoints struct.
To see whether our rate limiter is working, and since we set such low thresholds, we can test it just using our command-line tool.
First, restart the server (so the new code runs) by hitting Ctrl + C in the terminal window running the server. This signal will be trapped by our code, and an error will be sent down errChan, causing the program to quit. Once it has terminated, restart it:
go run main.go
Now, in another window, let's hash some passwords:
vaultcli hash bourgon
Repeat this command a few times–in most terminals, you can press the up arrow key and return. You'll notice that the first few requests succeed because it's within the limits, but if you get a little more aggressive and issue more than five requests in a second, you'll notice that we get errors:
$ vaultcli hash bourgon $2a$10$q3NTkjG0YFZhTG6gBU2WpenFmNzdN74oX0MDSTryiAqRXJ7RVw9sy $ vaultcli hash bourgon $2a$10$CdEEtxSDUyJEIFaykbMMl.EikxvV5921gs/./7If6VOdh2x0Q1oLW $ vaultcli hash bourgon $2a$10$1DSqQJJGCmVOptwIx6rrSOZwLlOhjHNC83OPVE8SdQ9q73Li5x2le $ vaultcli hash bourgon Invoke: rpc error: code = 2 desc = rate limit exceeded $ vaultcli hash bourgon Invoke: rpc error: code = 2 desc = rate limit exceeded $ vaultcli hash bourgon Invoke: rpc error: code = 2 desc = rate limit exceeded $ vaultcli hash bourgon $2a$10$kriTDXdyT6J4IrqZLwgBde663nLhoG3innhCNuf8H2nHf7kxnmSza
This shows that our rate limiter is working. We see errors until the token bucket fills back up, where our requests are fulfilled again.
Rather than returning an error (which is a pretty harsh response), perhaps we would prefer the server to just hold onto our request and fulfill it when it can-called throttling. For this case, Go kit provides the NewTokenBucketThrottler middleware.
Update the middleware code to use this middleware function instead:
hashEndpoint := vault.MakeHashEndpoint(srv)
{
hashEndpoint = ratelimitkit.NewTokenBucketThrottler(rlbucket,
time.Sleep)(hashEndpoint)
}
validateEndpoint := vault.MakeValidateEndpoint(srv)
{
validateEndpoint = ratelimitkit.NewTokenBucketThrottler(rlbucket,
time.Sleep)(validateEndpoint)
}
endpoints := vault.Endpoints{
HashEndpoint: hashEndpoint,
ValidateEndpoint: validateEndpoint,
}
The first argument to NewTokenBucketThrottler is the same endpoint as earlier, but now we have added a second argument of time.Sleep.
Note
Go kit allows us to customize the behavior by specifying what should happen when the delay needs to take place. In our case, we're passing time.Sleep, which is a function that will ask execution to pause for the specified amount of time. You could write your own function here if you wanted to do something different, but this works for now.
Now repeat the test from earlier, but this time, note that we never get an error-instead, the terminal will hang for a second until the request can be fulfilled.