github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/transport/tcp_response.go (about)

     1  package transport
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/volts-dev/volts/codec"
     7  	"github.com/volts-dev/volts/internal/body"
     8  )
     9  
    10  type RpcResponse struct {
    11  	body    *body.TBody //
    12  	sock    ISocket
    13  	Request *RpcRequest // request for this response
    14  	status  MessageStatusType
    15  }
    16  
    17  // 提供给Router的context使用
    18  func NewRpcResponse(ctx context.Context, req *RpcRequest, socket ISocket) *RpcResponse {
    19  	return &RpcResponse{
    20  		sock:    socket,
    21  		Request: req,
    22  		body:    body.New(codec.IdentifyCodec(req.Message.SerializeType())),
    23  	}
    24  }
    25  
    26  func (self *RpcResponse) Body() *body.TBody {
    27  	return self.body
    28  }
    29  
    30  // TODO 写状态
    31  func (self *RpcResponse) WriteHeader(code MessageStatusType) {
    32  	self.status = code
    33  }
    34  
    35  func (self *RpcResponse) Write(b []byte) (int, error) {
    36  	if self.Request.Message.IsOneway() {
    37  		return 0, nil // errors.New("This is one way request!")
    38  	}
    39  
    40  	msg := newMessage()
    41  	self.Request.Message.CloneTo(msg)
    42  	msg.SetMessageType(MT_RESPONSE)
    43  	if self.status == 0 {
    44  		self.status = StatusOK
    45  	}
    46  	msg.SetMessageStatusType(self.status)
    47  	msg.Payload = b
    48  	return len(b), self.sock.Send(msg)
    49  }
    50  
    51  // write data as stream
    52  func (self *RpcResponse) WriteStream(data interface{}) error {
    53  	if self.Request.Message.IsOneway() {
    54  		return nil // errors.New("This is one way request!")
    55  	}
    56  
    57  	msg := self.Request.Message
    58  	msg.SetMessageType(MT_RESPONSE)
    59  	if self.status == 0 {
    60  		self.status = StatusOK
    61  	}
    62  	msg.SetMessageStatusType(self.status)
    63  	_, err := self.body.Encode(data)
    64  	if err != nil {
    65  		return err
    66  	}
    67  
    68  	msg.Payload = self.Body().Data.Bytes()
    69  	return self.sock.Send(msg)
    70  }