github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/common/msgwriter/writer.go (about)

     1  package msgwriter
     2  
     3  import (
     4  	"github.com/nyan233/littlerpc/core/common/errorhandler"
     5  	"github.com/nyan233/littlerpc/core/common/inters"
     6  	"github.com/nyan233/littlerpc/core/common/transport"
     7  	"github.com/nyan233/littlerpc/core/middle/packer"
     8  	perror "github.com/nyan233/littlerpc/core/protocol/error"
     9  	"github.com/nyan233/littlerpc/core/protocol/message"
    10  	"sync"
    11  )
    12  
    13  const DefaultWriter = "lrpc-trait"
    14  
    15  // Writer 写入器的实现必须是线程安全的
    16  // 写入器的抽象与解析器不一样, 解析器要处理multi data & half package
    17  // 写入器使用的Conn API都是同步的, 所以不用处理half package, 写入器的设计
    18  // 本身就不能处理多份数据
    19  type Writer interface {
    20  	Write(arg Argument, header byte) perror.LErrorDesc
    21  	inters.Reset
    22  }
    23  
    24  type header interface {
    25  	Header() []byte
    26  }
    27  
    28  type Factory func(writers ...Writer) Writer
    29  
    30  type Argument struct {
    31  	Message *message.Message
    32  	Conn    transport.ConnAdapter
    33  	Encoder packer.Packer
    34  	// 用于统一内存复用的池, 类型是: *container.Slice[byte]
    35  	Pool *sync.Pool
    36  	// 不为nil时则说明Server开启了Debug模式
    37  	// 为true表示开启了Mux
    38  	OnDebug func([]byte, bool)
    39  	// 在消息发送完成时会调用
    40  	// 这个回调不应该用于调用插件的发送消息完成阶段, 一些Mux的实现可能会调用多次
    41  	// 这个回调, 如果想要唯一地检查应该去检查Write()的返回值
    42  	OnComplete func([]byte, perror.LErrorDesc)
    43  	EHandle    perror.LErrors
    44  }
    45  
    46  func encoder(arg Argument) perror.LErrorDesc {
    47  	// write body
    48  	if arg.Encoder.Scheme() != "text" {
    49  		bytes, err := arg.Encoder.EnPacket(arg.Message.Payloads())
    50  		if err != nil {
    51  			return arg.EHandle.LWarpErrorDesc(errorhandler.ErrServer, err.Error())
    52  		}
    53  		arg.Message.ReWritePayload(bytes)
    54  	}
    55  	return nil
    56  }
    57  
    58  var (
    59  	factoryCollection = make(map[string]Factory, 16)
    60  )
    61  
    62  func Register(scheme string, wf Factory) {
    63  	if wf == nil {
    64  		panic("Writer factory is nil")
    65  	}
    66  	if scheme == "" {
    67  		panic("Writer scheme is empty")
    68  	}
    69  	factoryCollection[scheme] = wf
    70  }
    71  
    72  func Get(scheme string) Factory {
    73  	return factoryCollection[scheme]
    74  }
    75  
    76  func init() {
    77  	Register(DefaultWriter, NewLRPCTrait)
    78  }