github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/server/stub.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"github.com/nyan233/littlerpc/core/container"
     7  	"github.com/nyan233/littlerpc/core/protocol/message"
     8  	"github.com/nyan233/littlerpc/core/utils/convert"
     9  	"unsafe"
    10  )
    11  
    12  // Stub 灵感来自arpc
    13  type Stub struct {
    14  	opt     *messageOpt
    15  	iter    *container.Iterator[[]byte]
    16  	reply   *message.Message
    17  	callErr error
    18  	context.Context
    19  }
    20  
    21  func (stub *Stub) setup() {
    22  	stub.iter = stub.opt.Message.PayloadsIterator()
    23  }
    24  
    25  func (stub *Stub) Read(p interface{}) error {
    26  	if !stub.iter.Next() {
    27  		return errors.New("read full")
    28  	}
    29  	bytes := stub.iter.Take()
    30  	switch p.(type) {
    31  	case *[]byte:
    32  		val := p.(*[]byte)
    33  		*val = append(*val, bytes...)
    34  	case *string:
    35  		slice := (*[]byte)(unsafe.Pointer(p.(*string)))
    36  		*slice = append(*slice, bytes...)
    37  	case nil:
    38  		break
    39  	default:
    40  		return stub.opt.Codec.Unmarshal(bytes, p)
    41  	}
    42  	return nil
    43  }
    44  
    45  var (
    46  	NullBytes = make([]byte, 0)
    47  )
    48  
    49  func (stub *Stub) Write(p interface{}) error {
    50  	switch p.(type) {
    51  	case nil:
    52  		stub.reply.AppendPayloads(NullBytes)
    53  	case *[]byte:
    54  		stub.reply.AppendPayloads(*p.(*[]byte))
    55  	case *string:
    56  		stub.reply.AppendPayloads(convert.StringToBytes(*p.(*string)))
    57  	case []byte:
    58  		stub.reply.AppendPayloads(p.([]byte))
    59  	case string:
    60  		stub.reply.AppendPayloads(convert.StringToBytes(p.(string)))
    61  	default:
    62  		bytes, err := stub.opt.Codec.Marshal(p)
    63  		if err != nil {
    64  			return err
    65  		}
    66  		stub.reply.AppendPayloads(bytes)
    67  	}
    68  	return nil
    69  }
    70  
    71  func (stub *Stub) WriteErr(err error) error {
    72  	stub.callErr = err
    73  	return nil
    74  }