grim/hgkeeper

Parents d3f0b014ec49
Children 4fc4d4c7aac6
Add a caching layer to the hgweb portion. This should take some strain off of mercurial anf our cpu quota
--- a/go.mod Mon Apr 13 23:14:50 2020 -0500
+++ b/go.mod Wed Apr 15 00:48:00 2020 -0500
@@ -6,6 +6,7 @@
github.com/casbin/casbin/v2 v2.0.2
github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568 // indirect
github.com/gliderlabs/ssh v0.2.2
+ github.com/hashicorp/golang-lru v0.5.4
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
github.com/sirupsen/logrus v1.4.1
github.com/smartystreets/goconvey v0.0.0-20190731233626-505e41936337 // indirect
--- a/go.sum Mon Apr 13 23:14:50 2020 -0500
+++ b/go.sum Wed Apr 15 00:48:00 2020 -0500
@@ -16,6 +16,8 @@
github.com/gliderlabs/ssh v0.2.2/go.mod h1:U7qILu1NlMHj9FlMhZLlkCdDnU1DBEAqr0aevW3Awn0=
github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8=
github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
+github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
+github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/hgweb/cache.go Wed Apr 15 00:48:00 2020 -0500
@@ -0,0 +1,96 @@
+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(),
+ })
+ }
+ })
+}
--- a/hgweb/hgweb.go Mon Apr 13 23:14:50 2020 -0500
+++ b/hgweb/hgweb.go Wed Apr 15 00:48:00 2020 -0500
@@ -17,11 +17,14 @@
listenAddr string
server *http.Server
cgiPath string
+
+ cacheSize int
}
-func NewServer(listenAddr string) (*Server, error) {
+func NewServer(listenAddr string, cacheSize int) (*Server, error) {
return &Server{
listenAddr: listenAddr,
+ cacheSize: cacheSize,
server: &http.Server{
Addr: listenAddr,
},
@@ -62,7 +65,12 @@
return err
}
- s.server.Handler = &cgi.Handler{Path: s.cgiPath}
+ cache, err := newCache(s.cacheSize)
+ if err != nil {
+ return err
+ }
+
+ s.server.Handler = cache.middleware(&cgi.Handler{Path: s.cgiPath})
log.Infof("http listening on %s", s.listenAddr)
--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/hgweb/responsewriter.go Wed Apr 15 00:48:00 2020 -0500
@@ -0,0 +1,68 @@
+package hgweb
+
+import (
+ "bytes"
+ "io"
+ "net/http"
+)
+
+type responseWriter struct {
+ w http.ResponseWriter
+ wroteHeader bool
+ statusCode int
+ written int64
+
+ multiWriter io.Writer
+ cachedBody *bytes.Buffer
+}
+
+var _ http.ResponseWriter = (*responseWriter)(nil)
+
+func newResponseWriter(w http.ResponseWriter) *responseWriter {
+ rw := &responseWriter{
+ w: w,
+ cachedBody: &bytes.Buffer{},
+ }
+
+ rw.multiWriter = io.MultiWriter(w, rw.cachedBody)
+
+ return rw
+}
+
+func (w *responseWriter) Header() http.Header {
+ return w.w.Header()
+}
+
+func (w *responseWriter) Write(data []byte) (int, error) {
+ if !w.wroteHeader {
+ w.wroteHeader = true
+ w.statusCode = http.StatusOK
+ }
+
+ l, err := w.multiWriter.Write(data)
+ w.written += int64(l)
+
+ return l, err
+}
+
+func (w *responseWriter) WriteHeader(statusCode int) {
+ if !w.wroteHeader {
+ w.wroteHeader = true
+
+ w.statusCode = statusCode
+ }
+
+ w.w.WriteHeader(w.statusCode)
+}
+
+func (w *responseWriter) StatusCode() int {
+ return w.statusCode
+}
+
+func (w *responseWriter) Written() int64 {
+ return w.written
+}
+
+func (w *responseWriter) Body() []byte {
+ return w.cachedBody.Bytes()
+}
--- a/serve/command.go Mon Apr 13 23:14:50 2020 -0500
+++ b/serve/command.go Wed Apr 15 00:48:00 2020 -0500
@@ -17,6 +17,7 @@
SSHAddr string `kong:"flag,name='ssh-listen-addr',env='HGK_SSH_LISTEN_ADDR',short='l',help='what address to listen on',default=':22222'"`
SSHHostKeysPath string `kong:"flag,name='ssh-host-keys-path',env='HGK_SSH_HOST_KEYS_PATH',short='H',help='the path where host keys are kept',default='host-keys'"`
HTTPAddr string `kong:"flag,name='http-listen-addr',env='HGK_HTTP_LISTEN_ADDR',help='what address the http server listens on',default=':8080'"`
+ CacheSize int `kong:"flag,name='cache-size',env='HGK_HTTP_CACHE_SIZE',help='number of pages to cache',default='1000'"`
}
func (c *Command) Run(g *globals.Globals) error {
@@ -36,7 +37,7 @@
}
defer ssh.Close()
- hgweb, err := hgweb.NewServer(c.HTTPAddr)
+ hgweb, err := hgweb.NewServer(c.HTTPAddr, c.CacheSize)
if err != nil {
return err
}