github.com/annwntech/go-micro/v2@v2.9.5/server/rpc_stream_test.go (about)

     1  package server
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"math/rand"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/golang/protobuf/proto"
    13  	"github.com/annwntech/go-micro/v2/codec/json"
    14  	protoCodec "github.com/annwntech/go-micro/v2/codec/proto"
    15  )
    16  
    17  // protoStruct implements proto.Message
    18  type protoStruct struct {
    19  	Payload string `protobuf:"bytes,1,opt,name=service,proto3" json:"service,omitempty"`
    20  }
    21  
    22  func (m *protoStruct) Reset()         { *m = protoStruct{} }
    23  func (m *protoStruct) String() string { return proto.CompactTextString(m) }
    24  func (*protoStruct) ProtoMessage()    {}
    25  
    26  // safeBuffer throws away everything and wont Read data back
    27  type safeBuffer struct {
    28  	sync.RWMutex
    29  	buf []byte
    30  	off int
    31  }
    32  
    33  func (b *safeBuffer) Write(p []byte) (n int, err error) {
    34  	if len(p) == 0 {
    35  		return 0, nil
    36  	}
    37  	// Cannot retain p, so we must copy it:
    38  	p2 := make([]byte, len(p))
    39  	copy(p2, p)
    40  	b.Lock()
    41  	b.buf = append(b.buf, p2...)
    42  	b.Unlock()
    43  	return len(p2), nil
    44  }
    45  
    46  func (b *safeBuffer) Read(p []byte) (n int, err error) {
    47  	if len(p) == 0 {
    48  		return 0, nil
    49  	}
    50  	b.RLock()
    51  	n = copy(p, b.buf[b.off:])
    52  	b.RUnlock()
    53  	if n == 0 {
    54  		return 0, io.EOF
    55  	}
    56  	b.off += n
    57  	return n, nil
    58  }
    59  
    60  func (b *safeBuffer) Close() error {
    61  	return nil
    62  }
    63  
    64  func TestRPCStream_Sequence(t *testing.T) {
    65  	buffer := new(bytes.Buffer)
    66  	rwc := readWriteCloser{
    67  		rbuf: buffer,
    68  		wbuf: buffer,
    69  	}
    70  	codec := json.NewCodec(&rwc)
    71  	streamServer := rpcStream{
    72  		codec: codec,
    73  		request: &rpcRequest{
    74  			codec: codec,
    75  		},
    76  	}
    77  
    78  	// Check if sequence is correct
    79  	for i := 0; i < 1000; i++ {
    80  		if err := streamServer.Send(fmt.Sprintf(`{"test":"value %d"}`, i)); err != nil {
    81  			t.Errorf("Unexpected Send error: %s", err)
    82  		}
    83  	}
    84  
    85  	for i := 0; i < 1000; i++ {
    86  		var msg string
    87  		if err := streamServer.Recv(&msg); err != nil {
    88  			t.Errorf("Unexpected Recv error: %s", err)
    89  		}
    90  		if msg != fmt.Sprintf(`{"test":"value %d"}`, i) {
    91  			t.Errorf("Unexpected msg: %s", msg)
    92  		}
    93  	}
    94  }
    95  
    96  func TestRPCStream_Concurrency(t *testing.T) {
    97  	buffer := new(safeBuffer)
    98  	codec := protoCodec.NewCodec(buffer)
    99  	streamServer := rpcStream{
   100  		codec: codec,
   101  		request: &rpcRequest{
   102  			codec: codec,
   103  		},
   104  	}
   105  
   106  	var wg sync.WaitGroup
   107  	// Check if race conditions happen
   108  	for i := 0; i < 10; i++ {
   109  		wg.Add(2)
   110  
   111  		go func() {
   112  			for i := 0; i < 50; i++ {
   113  				msg := protoStruct{Payload: "test"}
   114  				<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
   115  				if err := streamServer.Send(msg); err != nil {
   116  					t.Errorf("Unexpected Send error: %s", err)
   117  				}
   118  			}
   119  			wg.Done()
   120  		}()
   121  
   122  		go func() {
   123  			for i := 0; i < 50; i++ {
   124  				<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
   125  				if err := streamServer.Recv(&protoStruct{}); err != nil {
   126  					t.Errorf("Unexpected Recv error: %s", err)
   127  				}
   128  			}
   129  			wg.Done()
   130  		}()
   131  	}
   132  	wg.Wait()
   133  }