viney-shih/go-cache

package cache

import (
    "fmt"

    "github.com/klauspost/compress/s2"
    "github.com/vmihailenco/msgpack/v5"
)

// ref: https://github.com/go-redis/cache/blob/v8/cache.go

const (
    compressionThreshold = 64
    timeLen              = 4
)

const (
    noCompression = 0x0
    s2Compression = 0x1
)

// Marshal marshals value by msgpack + compress
func Marshal(value interface{}) ([]byte, error) {
    switch value := value.(type) {
    case nil:
        return nil, nil
    case []byte:
        return value, nil
    case string:
        return []byte(value), nil
    }

    b, err := msgpack.Marshal(value)
    if err != nil {
        return nil, err
    }

    return compress(b), nil
}

func compress(data []byte) []byte {
    if len(data) < compressionThreshold {
        n := len(data) + 1
        b := make([]byte, n, n+timeLen)
        copy(b, data)
        b[len(b)-1] = noCompression
        return b
    }

    n := s2.MaxEncodedLen(len(data)) + 1
    b := make([]byte, n, n+timeLen)
    b = s2.Encode(b, data)
    b = append(b, s2Compression)
    return b
}

// Unmarshal unmarshals binary with the compress + msgpack
func Unmarshal(b []byte, value interface{}) error {
    if len(b) == 0 {
        return nil
    }

    switch value := value.(type) {
    case nil:
        return nil
    case *[]byte:
        clone := make([]byte, len(b))
        copy(clone, b)
        *value = clone
        return nil
    case *string:
        *value = string(b)
        return nil
    }

    switch c := b[len(b)-1]; c {
    case noCompression:
        b = b[:len(b)-1]
    case s2Compression:
        b = b[:len(b)-1]

        var err error
        b, err = s2.Decode(nil, b)
        if err != nil {
            return err
        }
    default:
        return fmt.Errorf("unknown compression method: %x", c)
    }

    return msgpack.Unmarshal(b, value)
}