package cache import ( "bytes" "crypto/sha1" "io" "net/http" "net/url" "sync" "go-common/library/log" bm "go-common/library/net/http/blademaster" "go-common/library/net/http/blademaster/middleware/cache/store" proto "github.com/gogo/protobuf/proto" ) // consts for blademaster cache const ( _pagePrefix = "bm.page" ) // Page is used to cache common response type Page struct { Expire int32 pool sync.Pool } type cachedWriter struct { ctx *bm.Context response http.ResponseWriter store store.Store status int expire int32 key string } var _ http.ResponseWriter = &cachedWriter{} // NewPage will create a new page cache struct func NewPage(expire int32) *Page { pc := &Page{ Expire: expire, } pc.pool.New = func() interface{} { return &cachedWriter{} } return pc } // Key is used to identify response cache key in most key-value store func (p *Page) Key(ctx *bm.Context) string { url := ctx.Request.URL key := urlEscape(_pagePrefix, url.RequestURI()) return key } // Handler is used to execute cache service func (p *Page) Handler(store store.Store) bm.HandlerFunc { return func(ctx *bm.Context) { var ( resp *ResponseCache cached []byte err error ) key := p.Key(ctx) cached, err = store.Get(ctx, key) // if we did got the previous cache, // try to unmarshal it if err == nil && len(cached) > 0 { resp = new(ResponseCache) err = proto.Unmarshal(cached, resp) } // if we failed to fetch the cache or failed to parse cached data, // then consider try to cache this response if err != nil || resp == nil { writer := p.pool.Get().(*cachedWriter) writer.ctx = ctx writer.response = ctx.Writer writer.key = key writer.expire = p.Expire writer.store = store ctx.Writer = writer ctx.Next() p.pool.Put(writer) return } // write cached response headers := ctx.Writer.Header() for key, value := range resp.Header { headers[key] = value.Value } ctx.Writer.WriteHeader(int(resp.Status)) ctx.Writer.Write(resp.Data) ctx.Abort() } } func (w *cachedWriter) Header() http.Header { return w.response.Header() } func (w *cachedWriter) WriteHeader(code int) { w.status = int(code) w.response.WriteHeader(code) } func (w *cachedWriter) Write(data []byte) (size int, err error) { var ( origin []byte pdata []byte ) if size, err = w.response.Write(data); err != nil { return } store := w.store origin, err = store.Get(w.ctx, w.key) resp := new(ResponseCache) if err == nil || len(origin) > 0 { err1 := proto.Unmarshal(origin, resp) if err1 == nil { data = append(resp.Data, data...) } } resp.Status = int32(w.status) resp.Header = headerValues(w.Header()) resp.Data = data if pdata, err = proto.Marshal(resp); err != nil { // cannot happen log.Error("Failed to marshal response to protobuf: %v", err) return } if err = store.Set(w.ctx, w.key, pdata, w.expire); err != nil { log.Error("Failed to set response cache: %v", err) return } return } func headerValues(headers http.Header) map[string]*HeaderValue { result := make(map[string]*HeaderValue, len(headers)) for key, values := range headers { result[key] = &HeaderValue{ Value: values, } } return result } func urlEscape(prefix string, u string) string { key := url.QueryEscape(u) if len(key) > 200 { h := sha1.New() io.WriteString(h, u) key = string(h.Sum(nil)) } var buffer bytes.Buffer buffer.WriteString(prefix) buffer.WriteString(":") buffer.WriteString(key) return buffer.String() }