134 lines
2.7 KiB
Go
134 lines
2.7 KiB
Go
package mucp
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"math/rand"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/micro/go-micro/v3/codec/json"
|
|
protoCodec "github.com/micro/go-micro/v3/codec/proto"
|
|
)
|
|
|
|
// protoStruct implements proto.Message
|
|
type protoStruct struct {
|
|
Payload string `protobuf:"bytes,1,opt,name=service,proto3" json:"service,omitempty"`
|
|
}
|
|
|
|
func (m *protoStruct) Reset() { *m = protoStruct{} }
|
|
func (m *protoStruct) String() string { return proto.CompactTextString(m) }
|
|
func (*protoStruct) ProtoMessage() {}
|
|
|
|
// safeBuffer throws away everything and wont Read data back
|
|
type safeBuffer struct {
|
|
sync.RWMutex
|
|
buf []byte
|
|
off int
|
|
}
|
|
|
|
func (b *safeBuffer) Write(p []byte) (n int, err error) {
|
|
if len(p) == 0 {
|
|
return 0, nil
|
|
}
|
|
// Cannot retain p, so we must copy it:
|
|
p2 := make([]byte, len(p))
|
|
copy(p2, p)
|
|
b.Lock()
|
|
b.buf = append(b.buf, p2...)
|
|
b.Unlock()
|
|
return len(p2), nil
|
|
}
|
|
|
|
func (b *safeBuffer) Read(p []byte) (n int, err error) {
|
|
if len(p) == 0 {
|
|
return 0, nil
|
|
}
|
|
b.RLock()
|
|
n = copy(p, b.buf[b.off:])
|
|
b.RUnlock()
|
|
if n == 0 {
|
|
return 0, io.EOF
|
|
}
|
|
b.off += n
|
|
return n, nil
|
|
}
|
|
|
|
func (b *safeBuffer) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func TestRPCStream_Sequence(t *testing.T) {
|
|
buffer := new(bytes.Buffer)
|
|
rwc := readWriteCloser{
|
|
rbuf: buffer,
|
|
wbuf: buffer,
|
|
}
|
|
codec := json.NewCodec(&rwc)
|
|
streamServer := rpcStream{
|
|
codec: codec,
|
|
request: &rpcRequest{
|
|
codec: codec,
|
|
},
|
|
}
|
|
|
|
// Check if sequence is correct
|
|
for i := 0; i < 1000; i++ {
|
|
if err := streamServer.Send(fmt.Sprintf(`{"test":"value %d"}`, i)); err != nil {
|
|
t.Errorf("Unexpected Send error: %s", err)
|
|
}
|
|
}
|
|
|
|
for i := 0; i < 1000; i++ {
|
|
var msg string
|
|
if err := streamServer.Recv(&msg); err != nil {
|
|
t.Errorf("Unexpected Recv error: %s", err)
|
|
}
|
|
if msg != fmt.Sprintf(`{"test":"value %d"}`, i) {
|
|
t.Errorf("Unexpected msg: %s", msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRPCStream_Concurrency(t *testing.T) {
|
|
buffer := new(safeBuffer)
|
|
codec := protoCodec.NewCodec(buffer)
|
|
streamServer := rpcStream{
|
|
codec: codec,
|
|
request: &rpcRequest{
|
|
codec: codec,
|
|
},
|
|
}
|
|
|
|
var wg sync.WaitGroup
|
|
// Check if race conditions happen
|
|
for i := 0; i < 10; i++ {
|
|
wg.Add(2)
|
|
|
|
go func() {
|
|
for i := 0; i < 50; i++ {
|
|
msg := protoStruct{Payload: "test"}
|
|
<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
|
|
if err := streamServer.Send(msg); err != nil {
|
|
t.Errorf("Unexpected Send error: %s", err)
|
|
}
|
|
}
|
|
wg.Done()
|
|
}()
|
|
|
|
go func() {
|
|
for i := 0; i < 50; i++ {
|
|
<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
|
|
if err := streamServer.Recv(&protoStruct{}); err != nil {
|
|
t.Errorf("Unexpected Recv error: %s", err)
|
|
}
|
|
}
|
|
wg.Done()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
}
|