I had to capture metrics from a service I couldn’t directly access. Here’s how the go standard library solved my problem.
At determined.ai we enable teams of engineers to train deep learning models frictionlessly. With that said, how do models go from trained on our platform to available over the web for inference? One of our clients wanted to know exactly that. Enter the demo I was asked to build. In a single command, pull, deploy, and record response metrics from a model trained on our platform. Once the service comes online, visualize relevant metrics in Grafana.
We chose a TensorFlow model for our example. I pulled the checkpoint and wrapped it in a web service via Tensorflow Serving in an afternoon, which just left visualizing the model metrics. Turns out TensorFlow Serving doesn’t record the response metrics we were hoping to capture. I was running TensorFlow Serving as is, in it’s own process. For all intents and purposes I didn’t have access to the source for the service, so how could I record the metrics we needed?
Since the model is used over the web, clients could access the service via a reverse proxy. The proxy could then record metrics before responding to clients with their inference results. From there, exporting metrics to Prometheus and rendering some graphs in Grafana would be pretty simple.
”In computer networks , a reverse proxy is a type of proxy server that retrieves resources on behalf of a client from one or more servers. These resources are then returned to the client, appearing as if they originated from the proxy server itself. “
From Wikipedia
Essentially, a reverse proxy forwards traffic from a client to a set of servers behind the proxy. There are many applications for reverse proxies. Load balancing, TLS termination, and A/B testing are just a few. Reverse proxies are also useful for inserting instrumentation around an HTTP service without having to modify the service itself.
If you’d like to learn more about proxying I recommend checking out Introduction to modern network load balancing and proxying by Matt Klein. Matt is the creator of Envoy Proxy, a robust proxy server that powers service mesh tools like Istio. His post does a great job of outlining the approaches used by modern load balancers and proxies.
Go is one of my favorite programming languages for many reasons. The designers of the language focused on Simplicity, practicality, and performance. These considerations make Go a joy to use. The language shines with networking tasks. Part of the reason for this is the incredibly comprehensive standard library, which among other common implementations includes a reverse proxy 🤯.
Rolling your own proxy in go is as simple as https://gist.github.com/a9e65481bd6a3c44222f6967f5a60503
Yep, that’s it. Let’s dig in here. The httputil.NewSingleHostReverseProxy method returns a ReverseProxy struct containing the following method.
https://gist.github.com/a970c5cec1313ff8398faf84d1a1b493
All we need to do is configure the proxy and wire it up to a standard go HTTP server to have a working reverse proxy as shown below.
https://gist.github.com/2b2a25ffd030738225e64ea8119372b3
That’s it! This server can proxy HTTP requests and web socket connections. You’ll notice that I’ve configured the proxy.Director field. The ReverseProxy.Director is a function that modifies the incoming request before it is forwarded. The signature is as follows: https://gist.github.com/58171add61be4ea125b1472e6324a80a
A common use case for the director function is modifying request headers. One of the principles of the Go programming language is that types should have sane defaults and be immediately usable. Following this principle, the default director implementation returned by httputil.NewSingleHostReverseProxy takes care of setting the request Scheme, Host, and Path. I didn’t want to duplicate the code, so I wrapped that implementation. Note: I had to reset the req.Host field to handle HTTPS endpoints. I’ve also included an example of setting a request header via req.Header.Set which will override the header value with the value passed into the method.
Let’s extend our simple proxy to read and report metrics about the downstream service responses. To do this we’ll return to the httputil.ReverseProxy struct once more. It exposes a struct field ReverseProxy.ModifyResponse which gives us access to the HTTP response before it goes back to the client. https://gist.github.com/c0b91cffcff8cd6b8f46d6120f0f559a
Go implements HTTP bodies as io.Reader’s and, therefore, you may only read them once. If you would like to parse a request or response before forwarding it you will need to copy the body into a byte buffer and reset the body. An obvious drawback is that we buffer the entire response in memory without limit. This could lead to memory issues in production if you received a large response but for our use case this wasn’t an issue. Here’s a quick implementation to parse and reset the response body. https://gist.github.com/999ccdd2770a84a325148fe459b0fd78
With the request body problem solved, capturing metrics is simple. https://gist.github.com/e37942e4620ad2cbd8d2c6bcebd20968
And that’s it! The capture metrics function was pretty specific to my use case so I’ll leave it up to you to implement. I ended up using the Prometheus client library to predicted labels.
The full code for the metrics capturing proxy is as follows https://gist.github.com/d22cd037ed8387531feff45fcdc8b6a7
Go's standard library did all the heavy lifting necessary to get the demo up and running. From here you could extend the proxy to any number of use cases. Hopefully, you're as excited as I am about the practically of the go standard library. If you found this at all useful leave a comment below.