mirror of
https://github.com/hanzoai/psrpc.git
synced 2026-09-30 23:47:22 +00:00
add test middleware to simulate link latency (#78)
* add test middleware to simulate link latency * handle timer race * cleanup
This commit is contained in:
Paul Wells
+18
-21
@@ -44,28 +44,25 @@ func Proto() error {
|
||||
return err
|
||||
}
|
||||
|
||||
cmd := exec.Command(protoc,
|
||||
"--go_out", "internal",
|
||||
"--go_opt=paths=source_relative",
|
||||
"--plugin=go="+protocGoPath,
|
||||
"-I=./internal",
|
||||
"internal.proto",
|
||||
)
|
||||
mageutil.ConnectStd(cmd)
|
||||
if err = cmd.Run(); err != nil {
|
||||
return err
|
||||
protos := []struct {
|
||||
importPath, outputPath, filename string
|
||||
}{
|
||||
{"./internal", "internal", "internal.proto"},
|
||||
{"./protoc-gen-psrpc/options", "protoc-gen-psrpc/options", "options.proto"},
|
||||
{"./testutils", "testutils", "testutils.proto"},
|
||||
}
|
||||
|
||||
cmd = exec.Command(protoc,
|
||||
"--go_out", "protoc-gen-psrpc/options",
|
||||
"--go_opt=paths=source_relative",
|
||||
"--plugin=go="+protocGoPath,
|
||||
"-I=./protoc-gen-psrpc/options",
|
||||
"options.proto",
|
||||
)
|
||||
mageutil.ConnectStd(cmd)
|
||||
if err = cmd.Run(); err != nil {
|
||||
return err
|
||||
for _, p := range protos {
|
||||
cmd := exec.Command(protoc,
|
||||
"--go_out", p.outputPath,
|
||||
"--go_opt=paths=source_relative",
|
||||
"--plugin=go="+protocGoPath,
|
||||
"-I="+p.importPath,
|
||||
p.filename,
|
||||
)
|
||||
mageutil.ConnectStd(cmd)
|
||||
if err = cmd.Run(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
|
||||
@@ -147,6 +147,10 @@ func (m *multiRPC[ResponseType]) handleResponses(
|
||||
case <-timer.C:
|
||||
m.handler.Close()
|
||||
return
|
||||
|
||||
case <-ctx.Done():
|
||||
m.handler.Close()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,6 +39,14 @@ func WithSubscribeInterceptor(interceptor SubscribeInterceptor) TestBusOption {
|
||||
}
|
||||
}
|
||||
|
||||
func WithBusOptions(opts ...TestBusOption) TestBusOption {
|
||||
return func(o *bus.TestBusOpts) {
|
||||
for _, opt := range opts {
|
||||
opt(o)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func NewTestBus(next psrpc.MessageBus, opts ...TestBusOption) psrpc.MessageBus {
|
||||
return bus.NewTestBus(next, opts...)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
// Copyright 2023 LiveKit, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package testutils
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"context"
|
||||
"math"
|
||||
sync "sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/anypb"
|
||||
)
|
||||
|
||||
func WithLaggyBus(id string, latency latencyFunc) TestBusOption {
|
||||
return WithBusOptions(
|
||||
WithPublishInterceptor(func(next PublishHandler) PublishHandler {
|
||||
return func(ctx context.Context, channel string, msg proto.Message) error {
|
||||
a, err := anypb.New(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
b, err := proto.Marshal(a)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return next(ctx, channel, &LaggyMessage{
|
||||
Origin: id,
|
||||
SentAt: time.Now().UnixNano(),
|
||||
Body: b,
|
||||
})
|
||||
}
|
||||
}),
|
||||
WithSubscribeInterceptor(func(ctx context.Context, channel string, next ReadHandler) ReadHandler {
|
||||
l := newLaggySubscribeInterceptor()
|
||||
go l.Copy(ctx, id, latency, next)
|
||||
return l.Read
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
type latencyFunc func(a, b string) time.Duration
|
||||
|
||||
type laggySubscribeInterceptor struct {
|
||||
mu sync.Mutex
|
||||
rh readResultHeap
|
||||
t *time.Timer
|
||||
}
|
||||
|
||||
func newLaggySubscribeInterceptor() *laggySubscribeInterceptor {
|
||||
return &laggySubscribeInterceptor{
|
||||
t: time.NewTimer(math.MaxInt64),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *laggySubscribeInterceptor) pushRead(r *readResult) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
heap.Push(&l.rh, r)
|
||||
if l.rh.Peek() == r && l.t.Stop() {
|
||||
l.t.Reset(time.Until(r.time))
|
||||
}
|
||||
}
|
||||
|
||||
func (l *laggySubscribeInterceptor) popRead() ([]byte, bool) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
r := heap.Pop(&l.rh).(*readResult)
|
||||
if l.rh.Len() > 0 {
|
||||
l.t.Reset(time.Until(l.rh.Peek().time))
|
||||
} else if r.ok {
|
||||
l.t.Reset(math.MaxInt64)
|
||||
}
|
||||
return r.body, r.ok
|
||||
}
|
||||
|
||||
func (l *laggySubscribeInterceptor) Copy(
|
||||
ctx context.Context,
|
||||
id string,
|
||||
latency latencyFunc,
|
||||
read ReadHandler,
|
||||
) {
|
||||
for ctx.Err() == nil {
|
||||
b, ok := read()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
|
||||
a := &anypb.Any{}
|
||||
if proto.Unmarshal(b, a) != nil {
|
||||
break
|
||||
}
|
||||
|
||||
m := &LaggyMessage{}
|
||||
if a.UnmarshalTo(m) != nil {
|
||||
break
|
||||
}
|
||||
|
||||
l.pushRead(&readResult{
|
||||
time.Unix(0, m.SentAt).Add(latency(m.Origin, id)),
|
||||
m.Body,
|
||||
ok,
|
||||
})
|
||||
}
|
||||
|
||||
l.pushRead(&readResult{time: time.Now(), ok: false})
|
||||
}
|
||||
|
||||
func (l *laggySubscribeInterceptor) Read() (b []byte, open bool) {
|
||||
<-l.t.C
|
||||
return l.popRead()
|
||||
}
|
||||
|
||||
type readResult struct {
|
||||
time time.Time
|
||||
body []byte
|
||||
ok bool
|
||||
}
|
||||
|
||||
type readResultHeap struct {
|
||||
v []*readResult
|
||||
}
|
||||
|
||||
func (h *readResultHeap) Peek() *readResult { return h.v[0] }
|
||||
func (h *readResultHeap) Len() int { return len(h.v) }
|
||||
func (h *readResultHeap) Less(i, j int) bool { return h.v[i].time.Before(h.v[j].time) }
|
||||
func (h *readResultHeap) Swap(i, j int) { h.v[i], h.v[j] = h.v[j], h.v[i] }
|
||||
func (h *readResultHeap) Push(r any) { h.v = append(h.v, r.(*readResult)) }
|
||||
func (h *readResultHeap) Pop() any {
|
||||
r := h.v[len(h.v)-1]
|
||||
h.v[len(h.v)-1] = nil
|
||||
h.v = h.v[:len(h.v)-1]
|
||||
return r
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
// Copyright 2023 LiveKit, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Code generated by protoc-gen-go. DO NOT EDIT.
|
||||
// versions:
|
||||
// protoc-gen-go v1.31.0
|
||||
// protoc v4.23.4
|
||||
// source: testutils.proto
|
||||
|
||||
package testutils
|
||||
|
||||
import (
|
||||
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
|
||||
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
|
||||
reflect "reflect"
|
||||
sync "sync"
|
||||
)
|
||||
|
||||
const (
|
||||
// Verify that this generated code is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
|
||||
// Verify that runtime/protoimpl is sufficiently up-to-date.
|
||||
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
|
||||
)
|
||||
|
||||
type LaggyMessage struct {
|
||||
state protoimpl.MessageState
|
||||
sizeCache protoimpl.SizeCache
|
||||
unknownFields protoimpl.UnknownFields
|
||||
|
||||
Origin string `protobuf:"bytes,1,opt,name=origin,proto3" json:"origin,omitempty"`
|
||||
SentAt int64 `protobuf:"varint,2,opt,name=sent_at,json=sentAt,proto3" json:"sent_at,omitempty"`
|
||||
Body []byte `protobuf:"bytes,3,opt,name=body,proto3" json:"body,omitempty"`
|
||||
}
|
||||
|
||||
func (x *LaggyMessage) Reset() {
|
||||
*x = LaggyMessage{}
|
||||
if protoimpl.UnsafeEnabled {
|
||||
mi := &file_testutils_proto_msgTypes[0]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
}
|
||||
|
||||
func (x *LaggyMessage) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*LaggyMessage) ProtoMessage() {}
|
||||
|
||||
func (x *LaggyMessage) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_testutils_proto_msgTypes[0]
|
||||
if protoimpl.UnsafeEnabled && x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use LaggyMessage.ProtoReflect.Descriptor instead.
|
||||
func (*LaggyMessage) Descriptor() ([]byte, []int) {
|
||||
return file_testutils_proto_rawDescGZIP(), []int{0}
|
||||
}
|
||||
|
||||
func (x *LaggyMessage) GetOrigin() string {
|
||||
if x != nil {
|
||||
return x.Origin
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *LaggyMessage) GetSentAt() int64 {
|
||||
if x != nil {
|
||||
return x.SentAt
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *LaggyMessage) GetBody() []byte {
|
||||
if x != nil {
|
||||
return x.Body
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var File_testutils_proto protoreflect.FileDescriptor
|
||||
|
||||
var file_testutils_proto_rawDesc = []byte{
|
||||
0x0a, 0x0f, 0x74, 0x65, 0x73, 0x74, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74,
|
||||
0x6f, 0x12, 0x09, 0x74, 0x65, 0x73, 0x74, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x22, 0x53, 0x0a, 0x0c,
|
||||
0x4c, 0x61, 0x67, 0x67, 0x79, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x16, 0x0a, 0x06,
|
||||
0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x6f, 0x72,
|
||||
0x69, 0x67, 0x69, 0x6e, 0x12, 0x17, 0x0a, 0x07, 0x73, 0x65, 0x6e, 0x74, 0x5f, 0x61, 0x74, 0x18,
|
||||
0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x06, 0x73, 0x65, 0x6e, 0x74, 0x41, 0x74, 0x12, 0x12, 0x0a,
|
||||
0x04, 0x62, 0x6f, 0x64, 0x79, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x62, 0x6f, 0x64,
|
||||
0x79, 0x42, 0x24, 0x5a, 0x22, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f,
|
||||
0x6c, 0x69, 0x76, 0x65, 0x6b, 0x69, 0x74, 0x2f, 0x70, 0x73, 0x72, 0x70, 0x63, 0x2f, 0x74, 0x65,
|
||||
0x73, 0x74, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
|
||||
}
|
||||
|
||||
var (
|
||||
file_testutils_proto_rawDescOnce sync.Once
|
||||
file_testutils_proto_rawDescData = file_testutils_proto_rawDesc
|
||||
)
|
||||
|
||||
func file_testutils_proto_rawDescGZIP() []byte {
|
||||
file_testutils_proto_rawDescOnce.Do(func() {
|
||||
file_testutils_proto_rawDescData = protoimpl.X.CompressGZIP(file_testutils_proto_rawDescData)
|
||||
})
|
||||
return file_testutils_proto_rawDescData
|
||||
}
|
||||
|
||||
var file_testutils_proto_msgTypes = make([]protoimpl.MessageInfo, 1)
|
||||
var file_testutils_proto_goTypes = []interface{}{
|
||||
(*LaggyMessage)(nil), // 0: testutils.LaggyMessage
|
||||
}
|
||||
var file_testutils_proto_depIdxs = []int32{
|
||||
0, // [0:0] is the sub-list for method output_type
|
||||
0, // [0:0] is the sub-list for method input_type
|
||||
0, // [0:0] is the sub-list for extension type_name
|
||||
0, // [0:0] is the sub-list for extension extendee
|
||||
0, // [0:0] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_testutils_proto_init() }
|
||||
func file_testutils_proto_init() {
|
||||
if File_testutils_proto != nil {
|
||||
return
|
||||
}
|
||||
if !protoimpl.UnsafeEnabled {
|
||||
file_testutils_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
|
||||
switch v := v.(*LaggyMessage); i {
|
||||
case 0:
|
||||
return &v.state
|
||||
case 1:
|
||||
return &v.sizeCache
|
||||
case 2:
|
||||
return &v.unknownFields
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
File: protoimpl.DescBuilder{
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: file_testutils_proto_rawDesc,
|
||||
NumEnums: 0,
|
||||
NumMessages: 1,
|
||||
NumExtensions: 0,
|
||||
NumServices: 0,
|
||||
},
|
||||
GoTypes: file_testutils_proto_goTypes,
|
||||
DependencyIndexes: file_testutils_proto_depIdxs,
|
||||
MessageInfos: file_testutils_proto_msgTypes,
|
||||
}.Build()
|
||||
File_testutils_proto = out.File
|
||||
file_testutils_proto_rawDesc = nil
|
||||
file_testutils_proto_goTypes = nil
|
||||
file_testutils_proto_depIdxs = nil
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
// Copyright 2023 LiveKit, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package testutils;
|
||||
option go_package = "github.com/livekit/psrpc/testutils";
|
||||
|
||||
message LaggyMessage {
|
||||
string origin = 1;
|
||||
int64 sent_at = 2;
|
||||
bytes body = 3;
|
||||
}
|
||||
Reference in New Issue
Block a user