grim/hgkeeper

Add a caching layer to the hgweb portion. This should take some strain off of mercurial anf our cpu quota

2020-04-15, Gary Kramlich
410414a1fbe7
Add a caching layer to the hgweb portion. This should take some strain off of mercurial anf our cpu quota
package hgweb
import (
"net/http"
log "github.com/sirupsen/logrus"
"github.com/hashicorp/golang-lru"
)
type cache struct {
cache *lru.ARCCache
}
type cacheEntry struct {
headers http.Header
body []byte
}
func newCache(size int) (*cache, error) {
c, err := lru.NewARC(size)
if err != nil {
return nil, err
}
log.Debugf("created HTTP cache with size of %d", size)
return &cache{
cache: c,
}, nil
}
func (c *cache) middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// If this is not a GET request, we pass it straight through to
// mercurial.
if r.Method != http.MethodGet {
next.ServeHTTP(w, r)
return
}
// if the user didn't have a copy of this page, we check for a cached
// copy and if we have it, we return it back to them.
if r.Header.Get("If-None-Match") == "" {
if raw, ok := c.cache.Get(r.URL.String()); ok {
entry := raw.(cacheEntry)
newHeaders := w.Header()
for k, v := range entry.headers {
newHeaders[k] = v
}
w.Write(entry.body)
return
}
}
// At this point, the user did provide an etag or we don't have a
// cached copy of the page. So we pass the request on to mercurial,
// who will validate the etag if specified, or just do the work.
// create our response writer which caches the data
rw := newResponseWriter(w)
// send the request to mercurial
next.ServeHTTP(rw, r)
// if the status code is http.StatusOK, then either the users etag was
// expired or we didn't have this page cached. Regardless we then
// cache the page.
if rw.StatusCode() == http.StatusOK {
// create a http.Header skipping known bad headers
headers := http.Header{}
for k, v := range rw.Header() {
switch k {
case "Set-Cookie":
case "Authorization":
case "WWW-Authenticate":
case "Proxy-Authorization":
case "Proxy-Authenticate":
// these cases are all ignored.
default:
headers[k] = v
}
}
// add our entry to the cache
c.cache.Add(r.URL.String(), cacheEntry{
headers: headers,
body: rw.Body(),
})
}
})
}