live-gateway/douyu/stt/config.go

package stt

import (
	"github.com/modern-go/concurrent"
	"github.com/modern-go/reflect2"
	"io"
	"sync"
)

// Config customize how the API should behave.
// The API is created from Config by frozenConfig.
type Config struct {
	TagKey          string // 自定义 tagKey 默认 frozenConfig
	OnlyTaggedField bool   // 是否仅限 tag 标记字段
	CaseSensitive   bool   // key是否大小写敏感 默认true
}

// API the public interface of this package.
// Primary Marshal and Unmarshal.
type API interface {
	IteratorPool
	StreamPool

	Marshal(v any) ([]byte, error)
	Unmarshal(data []byte, v any) error
	RegisterExtension(extension Extension)
	DecoderOf(typ reflect2.Type) ValDecoder
	EncoderOf(typ reflect2.Type) ValEncoder
	//NewEncoder(writer io.Writer) *Encoder
	//NewDecoder(reader io.Reader) *Decoder
}

var ConfigDefault = Config{}.Froze()

type frozenConfig struct {
	configBeforeFrozen Config
	onlyTaggedField    bool
	caseSensitive      bool
	decoderCache       *concurrent.Map
	encoderCache       *concurrent.Map
	encoderExtension   Extension
	decoderExtension   Extension
	extraExtensions    []Extension
	streamPool         *sync.Pool // marshal stream
	iteratorPool       *sync.Pool // unmarshal iter pool
}

func (cfg *frozenConfig) initCache() {
	cfg.decoderCache = concurrent.NewMap()
	cfg.encoderCache = concurrent.NewMap()
}

func (cfg *frozenConfig) addDecoderToCache(cacheKey uintptr, decoder ValDecoder) {
	cfg.decoderCache.Store(cacheKey, decoder)
}

func (cfg *frozenConfig) addEncoderToCache(cacheKey uintptr, encoder ValEncoder) {
	cfg.encoderCache.Store(cacheKey, encoder)
}

func (cfg *frozenConfig) getDecoderFromCache(cacheKey uintptr) ValDecoder {
	decoder, found := cfg.decoderCache.Load(cacheKey)
	if found {
		return decoder.(ValDecoder)
	}
	return nil
}

func (cfg *frozenConfig) getEncoderFromCache(cacheKey uintptr) ValEncoder {
	encoder, found := cfg.encoderCache.Load(cacheKey)
	if found {
		return encoder.(ValEncoder)
	}
	return nil
}

var cfgCache = concurrent.NewMap()

func getFrozenConfigFromCache(cfg Config) *frozenConfig {
	obj, found := cfgCache.Load(cfg)
	if found {
		return obj.(*frozenConfig)
	}
	return nil
}

func addFrozenConfigToCache(cfg Config, frozenConfig *frozenConfig) {
	cfgCache.Store(cfg, frozenConfig)
}

///////////// Config

func (cfg Config) Froze() API {
	api := &frozenConfig{
		onlyTaggedField: cfg.OnlyTaggedField,
		caseSensitive:   cfg.CaseSensitive,
	}
	api.streamPool = &sync.Pool{
		New: func() interface{} {
			return NewStream(api, nil, 512)
		},
	}
	api.iteratorPool = &sync.Pool{
		New: func() interface{} {
			return NewIterator(api)
		},
	}
	api.initCache()
	encoderExtension := EncoderExtension{}
	decoderExtension := DecoderExtension{}

	api.encoderExtension = encoderExtension
	api.decoderExtension = decoderExtension
	api.configBeforeFrozen = cfg
	return api
}

func (cfg Config) frozeWithCacheReuse(extraExtensions []Extension) *frozenConfig {
	api := getFrozenConfigFromCache(cfg)
	if api != nil {
		return api
	}
	api = cfg.Froze().(*frozenConfig)
	for _, extension := range extraExtensions {
		api.RegisterExtension(extension)
	}
	addFrozenConfigToCache(cfg, api)
	return api
}

func (cfg *frozenConfig) getTagKey() string {
	tagKey := cfg.configBeforeFrozen.TagKey
	if tagKey == "" {
		return "stt"
	}
	return tagKey
}

func (cfg *frozenConfig) RegisterExtension(extension Extension) {
	cfg.extraExtensions = append(cfg.extraExtensions, extension)
	copied := cfg.configBeforeFrozen
	cfg.configBeforeFrozen = copied
}

func (cfg *frozenConfig) Marshal(v interface{}) ([]byte, error) {
	stream := cfg.BorrowStream(nil)
	defer cfg.ReturnStream(stream)
	stream.WriteVal(v)
	if stream.Error != nil {
		return nil, stream.Error
	}
	result := stream.Buffer()
	copied := make([]byte, len(result))
	copy(copied, result)
	return copied, nil
}

func (cfg *frozenConfig) Unmarshal(data []byte, v interface{}) error {
	// decrypt
	validData := cfg.getValidData(data)
	//iter := cfg.BorrowIterator(decryptBytes(validData, -1))
	iter := cfg.BorrowIterator(validData)
	defer cfg.ReturnIterator(iter)
	iter.ReadVal(v)
	c := iter.nextToken()
	if c == 0 {
		if iter.Error == io.EOF {
			return nil
		}
		return iter.Error
	}
	iter.ReportError("Unmarshal", "there are bytes left after unmarshal")
	return iter.Error
}

func (cfg *frozenConfig) getValidData(data []byte) []byte {
	c := data[len(data)-1]
	if c != '/' {
		data = append(data, '/')
	}
	return data
}