github.com/sandwich-go/boost@v1.3.29/xencoding/codec.go (about)

     1  package xencoding
     2  
     3  import (
     4  	"context"
     5  	"github.com/sandwich-go/boost/xpanic"
     6  	"sort"
     7  	"sync"
     8  )
     9  
    10  type codecKeyType struct{}
    11  
    12  // for ark context
    13  func (*codecKeyType) String() string { return "encoding2-codec—key" }
    14  
    15  var keyForContext = codecKeyType(struct{}{})
    16  
    17  // WithContext 将 Codec 存放在 context.Context 中
    18  func WithContext(ctx context.Context, c Codec) context.Context {
    19  	return context.WithValue(ctx, keyForContext, c)
    20  }
    21  
    22  // FromContext 从 context.Context 中 获取 Codec
    23  func FromContext(ctx context.Context) Codec {
    24  	c := ctx.Value(keyForContext)
    25  	if c == nil {
    26  		c = ctx.Value(keyForContext.String())
    27  	}
    28  	if c == nil {
    29  		return nil
    30  	}
    31  	return c.(Codec)
    32  }
    33  
    34  // Codec defines the interface link uses to encode and decode messages.
    35  type Codec interface {
    36  	// Marshal returns the wire format of v.
    37  	Marshal(context.Context, interface{}) ([]byte, error)
    38  	// Unmarshal parses the wire format into v.
    39  	Unmarshal(context.Context, []byte, interface{}) error
    40  	// Name String returns the name of the Codec implementation.
    41  	Name() string
    42  }
    43  
    44  var (
    45  	mu     sync.RWMutex
    46  	codecs = make(map[string]Codec)
    47  )
    48  
    49  // RegisterCodec 注册 Codec
    50  // 可以注册自定义的 Codec
    51  func RegisterCodec(c Codec) {
    52  	mu.Lock()
    53  	defer mu.Unlock()
    54  
    55  	xpanic.WhenTrue(c == nil, "cannot register a nil Codec")
    56  	name := c.Name()
    57  	xpanic.WhenTrue(len(name) == 0, "cannot register Codec with empty string result for Name()")
    58  	_, dup := codecs[name]
    59  	xpanic.WhenTrue(dup, "register called twice for codec %s", name)
    60  	codecs[name] = c
    61  }
    62  
    63  // Codecs 获取所有 Codec 的名称
    64  func Codecs() []string {
    65  	mu.RLock()
    66  	defer mu.RUnlock()
    67  
    68  	list := make([]string, 0, len(codecs))
    69  	for name := range codecs {
    70  		list = append(list, name)
    71  	}
    72  	sort.Strings(list)
    73  	return list
    74  }
    75  
    76  // GetCodec 通过名称来获取注册过的 Codec
    77  func GetCodec(name string) Codec {
    78  	mu.RLock()
    79  	defer mu.RUnlock()
    80  
    81  	return codecs[name]
    82  }