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

     1  package msgwriter
     2  
     3  import (
     4  	perror "github.com/nyan233/littlerpc/core/protocol/error"
     5  )
     6  
     7  type lRPCTrait struct {
     8  	handlers map[byte]Writer
     9  }
    10  
    11  func NewLRPCTrait(writers ...Writer) Writer {
    12  	trait := &lRPCTrait{
    13  		handlers: make(map[byte]Writer, 16),
    14  	}
    15  	handlers := []Writer{NewLRPCMux(), NewJsonRPC2(), NewLRPCNoMux()}
    16  	handlers = append(handlers, writers...)
    17  	for _, handler := range handlers {
    18  		headerI, ok := handler.(header)
    19  		if ok {
    20  			for _, v := range headerI.Header() {
    21  				trait.handlers[v] = handler
    22  			}
    23  		}
    24  	}
    25  	return trait
    26  }
    27  
    28  func (l *lRPCTrait) Write(arg Argument, header byte) perror.LErrorDesc {
    29  	writer := l.handlers[header]
    30  	return writer.Write(arg, header)
    31  }
    32  
    33  func (l *lRPCTrait) Reset() {
    34  	return
    35  }