GoBlog/cache.go

127 lines
3.0 KiB
Go

package main
import (
"io/ioutil"
"net/http"
"net/http/httptest"
"sync"
"time"
"golang.org/x/sync/singleflight"
)
var (
cacheGroup singleflight.Group
cacheMap = map[string]*cacheItem{}
cacheMutex = &sync.RWMutex{}
)
func initCache() {
go func() {
for {
// GC the entries every 60 seconds
time.Sleep(60 * time.Second)
cacheMutex.Lock()
for key, item := range cacheMap {
if item.expired() {
delete(cacheMap, key)
}
}
cacheMutex.Unlock()
}
}()
}
func cacheMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if appConfig.Cache.Enable &&
// check method
(r.Method == http.MethodGet || r.Method == http.MethodHead) &&
// check bypass query
!(r.URL.Query().Get("cache") == "0") {
key := cacheKey(r)
// Get cache or render it
cacheInterface, _, _ := cacheGroup.Do(key, func() (interface{}, error) {
return getCache(key, next, r), nil
})
cache := cacheInterface.(*cacheItem)
// log.Println(string(cache.body))
cacheTimeString := time.Unix(cache.creationTime, 0).Format(time.RFC1123)
expiresTimeString := time.Unix(cache.creationTime+appConfig.Cache.Expiration, 0).Format(time.RFC1123)
// check conditional request
ifModifiedSinceHeader := r.Header.Get("If-Modified-Since")
if ifModifiedSinceHeader != "" && ifModifiedSinceHeader == cacheTimeString {
setCacheHeaders(w, cacheTimeString, expiresTimeString)
// send 304
w.WriteHeader(http.StatusNotModified)
return
}
// copy cached headers
for k, v := range cache.header {
w.Header()[k] = v
}
setCacheHeaders(w, cacheTimeString, expiresTimeString)
// set status code
w.WriteHeader(cache.code)
// write cached body
_, _ = w.Write(cache.body)
return
}
next.ServeHTTP(w, r)
return
})
}
func cacheKey(r *http.Request) string {
return r.URL.String()
}
func setCacheHeaders(w http.ResponseWriter, cacheTimeString string, expiresTimeString string) {
w.Header().Set("Cache-Control", "public")
w.Header().Set("Last-Modified", cacheTimeString)
w.Header().Set("Expires", expiresTimeString)
}
type cacheItem struct {
creationTime int64
code int
header http.Header
body []byte
}
func (c *cacheItem) expired() bool {
return c.creationTime < time.Now().Unix()-appConfig.Cache.Expiration
}
func getCache(key string, next http.Handler, r *http.Request) *cacheItem {
cacheMutex.RLock()
item, ok := cacheMap[key]
cacheMutex.RUnlock()
if !ok || item.expired() {
// No cache available
// Record request
recorder := httptest.NewRecorder()
next.ServeHTTP(recorder, r)
// Cache values from recorder
result := recorder.Result()
body, _ := ioutil.ReadAll(result.Body)
item = &cacheItem{
creationTime: time.Now().Unix(),
code: result.StatusCode,
header: result.Header,
body: body,
}
// Save cache
cacheMutex.Lock()
cacheMap[key] = item
cacheMutex.Unlock()
}
return item
}
func purgeCache() {
cacheMutex.Lock()
cacheMap = map[string]*cacheItem{}
cacheMutex.Unlock()
}