Implement cache and logger middlewares for gemini

7 files changed, 189 insertions, 23 deletions

diff --git a/src/cache/cache.go b/src/cache/cache.go
new file mode 100644
index 0000000..ade286a
--- /dev/null
+++ b/src/cache/cache.go
@@ -0,0 +1,49 @@
+// Package cache implements a simple LRU cache which can be completely purged
+// whenever needed.
+package cache
+
+import (
+	lru "github.com/hashicorp/golang-lru"
+)
+
+// Cache describes an in-memory cache for arbitrary objects, which has a Purge
+// method for clearing everything at once.
+type Cache interface {
+	Get(key string) interface{}
+	Set(key string, value interface{})
+	Purge()
+}
+
+type cache struct {
+	cache *lru.Cache
+}
+
+// New instantiates and returns a new Cache which can hold up to the given
+// number of entries.
+func New(size int) Cache {
+	c, err := lru.New(size)
+
+	// instantiating the lru cache can't realistically fail
+	if err != nil {
+		panic(err)
+	}
+
+	return &cache{c}
+}
+
+func (c *cache) Get(key string) interface{} {
+	value, ok := c.cache.Get(key)
+	if !ok {
+		return nil
+	}
+
+	return value
+}
+
+func (c *cache) Set(key string, value interface{}) {
+	c.cache.Add(key, value)
+}
+
+func (c *cache) Purge() {
+	c.cache.Purge()
+}
diff --git a/src/cmd/mediocre-blog/main.go b/src/cmd/mediocre-blog/main.go
index ff8e478..d6a6b56 100644
--- a/src/cmd/mediocre-blog/main.go
+++ b/src/cmd/mediocre-blog/main.go
@@ -7,6 +7,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
 	cfgpkg "github.com/mediocregopher/blog.mediocregopher.com/srv/cfg"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/gmi"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/http"
@@ -102,7 +103,10 @@ func main() {
 	postAssetStore := post.NewAssetStore(postSQLDB)
 	postDraftStore := post.NewDraftStore(postSQLDB)
 
+	cache := cache.New(5000)
+
 	httpParams.Logger = logger.WithNamespace("http")
+	httpParams.Cache = cache
 	httpParams.PowManager = powMgr
 	httpParams.PostStore = postStore
 	httpParams.PostAssetStore = postAssetStore
@@ -124,6 +128,7 @@ func main() {
 	}()
 
 	gmiParams.Logger = logger.WithNamespace("gmi")
+	gmiParams.Cache = cache
 	gmiParams.PostStore = postStore
 	gmiParams.PostAssetStore = postAssetStore
 	gmiParams.HTTPPublicURL = httpParams.PublicURL
diff --git a/src/gmi/cache.go b/src/gmi/cache.go
new file mode 100644
index 0000000..0f2e975
--- /dev/null
+++ b/src/gmi/cache.go
@@ -0,0 +1,87 @@
+package gmi
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"sync"
+
+	"git.sr.ht/~adnano/go-gemini"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
+)
+
+type cacheRW struct {
+	gemini.ResponseWriter
+	status    gemini.Status
+	mediaType string
+	buf       []byte
+}
+
+func (c *cacheRW) SetMediaType(mediaType string) {
+	c.mediaType = mediaType
+	c.ResponseWriter.SetMediaType(mediaType)
+}
+
+func (c *cacheRW) WriteHeader(status gemini.Status, meta string) {
+	c.status = status
+	c.ResponseWriter.WriteHeader(status, meta)
+}
+
+func (c *cacheRW) Write(b []byte) (int, error) {
+	c.buf = append(c.buf, b...)
+	return c.ResponseWriter.Write(b)
+}
+
+func cacheMiddleware(cache cache.Cache) func(h gemini.Handler) gemini.Handler {
+
+	type entry struct {
+		mediaType string
+		body      []byte
+	}
+
+	pool := sync.Pool{
+		New: func() interface{} { return new(bytes.Reader) },
+	}
+
+	return func(h gemini.Handler) gemini.Handler {
+		return gemini.HandlerFunc(func(
+			ctx context.Context,
+			rw gemini.ResponseWriter,
+			r *gemini.Request,
+		) {
+
+			id := r.URL.String()
+
+			if value := cache.Get(id); value != nil {
+
+				entry := value.(entry)
+
+				if entry.mediaType != "" {
+					rw.SetMediaType(entry.mediaType)
+				}
+
+				reader := pool.Get().(*bytes.Reader)
+				defer pool.Put(reader)
+
+				reader.Reset(entry.body)
+
+				io.Copy(rw, reader)
+				return
+			}
+
+			cacheRW := &cacheRW{
+				ResponseWriter: rw,
+				status:         gemini.StatusSuccess,
+			}
+
+			h.ServeGemini(ctx, cacheRW, r)
+
+			if cacheRW.status == gemini.StatusSuccess {
+				cache.Set(id, entry{
+					mediaType: cacheRW.mediaType,
+					body:      cacheRW.buf,
+				})
+			}
+		})
+	}
+}
diff --git a/src/gmi/gmi.go b/src/gmi/gmi.go
index 6835ea0..c2a32aa 100644
--- a/src/gmi/gmi.go
+++ b/src/gmi/gmi.go
@@ -13,6 +13,7 @@ import (
 
 	"git.sr.ht/~adnano/go-gemini"
 	"git.sr.ht/~adnano/go-gemini/certificate"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/cfg"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/post"
 	"github.com/mediocregopher/mediocre-go-lib/v2/mctx"
@@ -23,6 +24,7 @@ import (
 // unless otherwise noted.
 type Params struct {
 	Logger *mlog.Logger
+	Cache  cache.Cache
 
 	PostStore      post.Store
 	PostAssetStore post.AssetStore
@@ -108,10 +110,11 @@ func New(params Params) (API, error) {
 
 	go func() {
 
-		ctx := mctx.WithAnnotator(context.Background(), &a.params)
+		err := a.srv.ListenAndServe(context.Background())
 
-		err := a.srv.ListenAndServe(ctx)
 		if err != nil && !errors.Is(err, context.Canceled) {
+
+			ctx := mctx.WithAnnotator(context.Background(), &a.params)
 			a.params.Logger.Fatal(ctx, "serving gemini server", err)
 		}
 	}()
@@ -123,6 +126,26 @@ func (a *api) Shutdown(ctx context.Context) error {
 	return a.srv.Shutdown(ctx)
 }
 
+func (a *api) logReqMiddleware(h gemini.Handler) gemini.Handler {
+
+	type logCtxKey string
+
+	return gemini.HandlerFunc(func(
+		ctx context.Context,
+		rw gemini.ResponseWriter,
+		r *gemini.Request,
+	) {
+
+		ctx = mctx.Annotate(ctx,
+			logCtxKey("url"), r.URL.String(),
+		)
+
+		h.ServeGemini(ctx, rw, r)
+
+		a.params.Logger.Info(ctx, "handled gemini request")
+	})
+}
+
 func indexMiddleware(h gemini.Handler) gemini.Handler {
 
 	return gemini.HandlerFunc(func(
@@ -212,8 +235,8 @@ func (a *api) handler() (gemini.Handler, error) {
 	h = mux
 	h = indexMiddleware(h)
 
-	// TODO logging
-	// TODO caching
+	h = a.logReqMiddleware(h)
+	h = cacheMiddleware(a.params.Cache)(h)
 
 	return h, nil
 }
diff --git a/src/http/apiutil/apiutil.go b/src/http/apiutil/apiutil.go
index 1fbadea..c59104c 100644
--- a/src/http/apiutil/apiutil.go
+++ b/src/http/apiutil/apiutil.go
@@ -16,6 +16,9 @@ import (
 	"github.com/mediocregopher/mediocre-go-lib/v2/mlog"
 )
 
+// TODO I don't think Set/GetRequestLogger are necessary? Seems sufficient to
+// just annotate the request's context
+
 type loggerCtxKey int
 
 // SetRequestLogger sets the given Logger onto the given Request's Context,
diff --git a/src/http/http.go b/src/http/http.go
index 98cdde3..d51671d 100644
--- a/src/http/http.go
+++ b/src/http/http.go
@@ -15,7 +15,7 @@ import (
 	"strings"
 	"time"
 
-	lru "github.com/hashicorp/golang-lru"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/cfg"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/http/apiutil"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/mailinglist"
@@ -33,6 +33,7 @@ var staticFS embed.FS
 type Params struct {
 	Logger     *mlog.Logger
 	PowManager pow.Manager
+	Cache      cache.Cache
 
 	PostStore      post.Store
 	PostAssetStore post.AssetStore
@@ -190,13 +191,6 @@ func (a *api) apiHandler() http.Handler {
 
 func (a *api) blogHandler() http.Handler {
 
-	cache, err := lru.New(5000)
-
-	// instantiating the lru cache can't realistically fail
-	if err != nil {
-		panic(err)
-	}
-
 	mux := http.NewServeMux()
 
 	mux.Handle("/posts/", http.StripPrefix("/posts",
@@ -244,11 +238,11 @@ func (a *api) blogHandler() http.Handler {
 
 	readOnlyMiddlewares := []middleware{
 		logReqMiddleware, // only log GETs on cache miss
-		cacheMiddleware(cache),
+		cacheMiddleware(a.params.Cache, a.params.PublicURL),
 	}
 
 	readWriteMiddlewares := []middleware{
-		purgeCacheOnOKMiddleware(cache),
+		purgeCacheOnOKMiddleware(a.params.Cache),
 		authMiddleware(a.auther),
 	}
 
diff --git a/src/http/middleware.go b/src/http/middleware.go
index b82fc29..a21511f 100644
--- a/src/http/middleware.go
+++ b/src/http/middleware.go
@@ -4,11 +4,12 @@ import (
 	"bytes"
 	"net"
 	"net/http"
+	"net/url"
 	"path/filepath"
 	"sync"
 	"time"
 
-	lru "github.com/hashicorp/golang-lru"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/http/apiutil"
 	"github.com/mediocregopher/mediocre-go-lib/v2/mctx"
 	"github.com/mediocregopher/mediocre-go-lib/v2/mlog"
@@ -86,6 +87,9 @@ func (rw *wrappedResponseWriter) WriteHeader(statusCode int) {
 }
 
 func logReqMiddleware(h http.Handler) http.Handler {
+
+	type logCtxKey string
+
 	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 
 		wrw := newWrappedResponseWriter(rw)
@@ -94,8 +98,6 @@ func logReqMiddleware(h http.Handler) http.Handler {
 		h.ServeHTTP(wrw, r)
 		took := time.Since(started)
 
-		type logCtxKey string
-
 		ctx := r.Context()
 		ctx = mctx.Annotate(ctx,
 			logCtxKey("took"), took.String(),
@@ -139,7 +141,7 @@ func (rw *cacheResponseWriter) Write(b []byte) (int, error) {
 	return rw.wrappedResponseWriter.Write(b)
 }
 
-func cacheMiddleware(cache *lru.Cache) middleware {
+func cacheMiddleware(cache cache.Cache, publicURL *url.URL) middleware {
 
 	type entry struct {
 		body      []byte
@@ -153,11 +155,14 @@ func cacheMiddleware(cache *lru.Cache) middleware {
 	return func(h http.Handler) http.Handler {
 		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 
-			id := r.URL.RequestURI()
+			// r.URL doesn't have Scheme or Host populated, better to add the
+			// public url to the key to make sure there's no possiblity of
+			// collision with other protocols using the cache.
+			id := publicURL.String() + "|" + r.URL.String()
 
-			if val, ok := cache.Get(id); ok {
+			if value := cache.Get(id); value != nil {
 
-				entry := val.(entry)
+				entry := value.(entry)
 
 				reader := pool.Get().(*bytes.Reader)
 				defer pool.Put(reader)
@@ -174,7 +179,7 @@ func cacheMiddleware(cache *lru.Cache) middleware {
 			h.ServeHTTP(cacheRW, r)
 
 			if cacheRW.statusCode == 200 {
-				cache.Add(id, entry{
+				cache.Set(id, entry{
 					body:      cacheRW.buf.Bytes(),
 					createdAt: time.Now(),
 				})
@@ -183,7 +188,7 @@ func cacheMiddleware(cache *lru.Cache) middleware {
 	}
 }
 
-func purgeCacheOnOKMiddleware(cache *lru.Cache) middleware {
+func purgeCacheOnOKMiddleware(cache cache.Cache) middleware {
 	return func(h http.Handler) http.Handler {
 		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {