kratos-utils/transport/http/context/stream_context.go

package context

import (
	"context"
	"encoding/json"
)

type StreamContext interface {
	context.Context
	Serialize(any) ([]byte, error)
}

type streamContext struct {
	context.Context
}

type Serialize func(any) ([]byte, error)

var _ StreamContext = (*streamContext)(nil)

var (
	// serialize 是全局的自定义序列化器
	serialize Serialize
)

func (s *streamContext) Serialize(v any) ([]byte, error) {
	if serialize != nil {
		return serialize(v)
	}
	return json.Marshal(v)
}

func NewStreamContext(ctx context.Context) StreamContext {
	return &streamContext{
		Context: ctx,
	}
}

func SetSerializer(s Serialize) {
	serialize = s
}