METANOIA/metadata/cache.go

package metadata

import (
	"bytes"
	"encoding/json"
	"errors"
	"github.com/dgraph-io/badger/v3"
	badgerOptions "github.com/dgraph-io/badger/v3/options"
	"github.com/minio/sha256-simd"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
	"time"
)

type BadgerCacheStore struct {
	handle       *badger.DB
	gcTicker     *time.Ticker
	closeChannel chan bool
}

func NewBadgerCacheStore(path string) (*BadgerCacheStore, error) {
	options := badger.DefaultOptions(path)
	options.SyncWrites = false
	options.NumVersionsToKeep = 1
	options.Compression = badgerOptions.ZSTD
	options.ZSTDCompressionLevel = 1
	db, err := badger.Open(options)
	if err != nil {
		return nil, err
	}

	store := &BadgerCacheStore{
		handle:       db,
		gcTicker:     time.NewTicker(time.Minute * 5),
		closeChannel: make(chan bool),
	}
	go func() {
		defer store.gcTicker.Stop()

		for {
			select {
			case <-store.gcTicker.C:
				for store.handle.RunValueLogGC(0.6) == nil {

				}
			case <-store.closeChannel:
				return
			}
		}
	}()

	runtime.SetFinalizer(store, func(s *BadgerCacheStore) {
		s.closeChannel <- true
		s.handle.Close()
	})

	return store, nil
}

func getRequestKey(r *http.Request) (out []byte) {
	key := r.Method + ":" + r.URL.String() + ":"
	var headers []string
	for k, v := range r.Header {
		hk := strings.ToLower(k)
		if hk == "user-agent" || hk == "x-fetched-at" {
			continue
		}
		headers = append(headers, hk+":"+strings.Join(v, ","))
	}
	sort.SliceStable(headers, func(i, j int) bool {
		return strings.Compare(headers[i], headers[j]) < 0
	})
	key += strings.Join(headers, ";")

	hasher := sha256.New()
	hasher.Write([]byte(key))

	out = hasher.Sum(out)

	return
}

type encodedRequest struct {
	Method  string              `json:"method"`
	URL     string              `json:"url"`
	Headers map[string][]string `json:"headers"`
}

func encodeRequest(r *http.Request) (out []byte) {

	value := encodedRequest{
		Method:  r.Method,
		URL:     r.URL.String(),
		Headers: r.Header,
	}

	out, _ = json.Marshal(value)
	return
}

type encodedResponse struct {
	Request    encodedRequest      `json:"request"`
	Headers    map[string][]string `json:"headers"`
	Status     string              `json:"status"`
	StatusCode int                 `json:"status_code"`
	Proto      string              `json:"proto"`
	Body       []byte              `json:"body"`
}

func encodeResponse(r *http.Response) (out []byte) {

	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		return nil
	}

	value := encodedResponse{
		Request: encodedRequest{
			Method:  r.Request.Method,
			URL:     r.Request.URL.String(),
			Headers: r.Request.Header,
		},
		Headers:    r.Header,
		Status:     r.Status,
		StatusCode: r.StatusCode,
		Proto:      r.Proto,
		Body:       body,
	}

	out, _ = json.Marshal(value)
	return
}

func (s *BadgerCacheStore) Get(request *http.Request) (response *http.Response, err error) {
	key := getRequestKey(request)

	err = s.handle.View(func(txn *badger.Txn) error {
		item, err := txn.Get(key)
		if err != nil {
			return err
		}

		err = item.Value(func(val []byte) error {
			value := &encodedResponse{}

			err := json.Unmarshal(val, value)
			if err != nil {
				return err
			}

			requestUri, _ := url.Parse(value.Request.URL)
			response = &http.Response{
				Request: &http.Request{
					Method: value.Request.Method,
					URL:    requestUri,
					Header: value.Request.Headers,
				},
				Header:     value.Headers,
				Status:     value.Status,
				StatusCode: value.StatusCode,
				Proto:      value.Proto,
				Body:       io.NopCloser(bytes.NewReader(value.Body)),
			}

			return nil
		})
		return err
	})

	return
}

func (s *BadgerCacheStore) Set(request *http.Request, response *http.Response) (*http.Response, error) {
	defer response.Body.Close()
	key := getRequestKey(request)
	byteValue := encodeResponse(response)
	if len(byteValue) == 0 {
		return nil, errors.New("could not encode response")
	}

	err := s.handle.Update(func(txn *badger.Txn) error {
		return txn.Set(key, byteValue)
	})
	if err != nil {
		return nil, err
	}

	return s.Get(request)
}