hanzoai/psrpc
2
0
Fork 0
mirror of https://github.com/hanzoai/psrpc.git synced 2026-09-30 23:47:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Affinity context (#71)

Browse Source 
* add context to affinity

* update version
This commit is contained in:
David Colburn
2023-10-23 16:29:31 -07:00
committed by GitHub
parent dc1506aa6a
commit f9ce87b4f2
5 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/test/my_service/my_service_test.go
+1 -1
View File
@@ -272,7 +272,7 @@ func (s *MyService) IntensiveRPC(_ context.Context, _ *MyRequest) (*MyResponse,
return &MyResponse{}, nil
}
func (s *MyService) IntensiveRPCAffinity(_ *MyRequest) float32 {
func (s *MyService) IntensiveRPCAffinity(_ context.Context, _ *MyRequest) float32 {
s.Lock()
s.counts["IntensiveRPCAffinity"]++
s.Unlock()
pkg/server/rpc.go
+2 -2
View File
@@ -30,7 +30,7 @@ import (
"github.com/livekit/psrpc/pkg/metadata"
)
type AffinityFunc[RequestType proto.Message] func(RequestType) float32
type AffinityFunc[RequestType proto.Message] func(context.Context, RequestType) float32
type rpcHandlerImpl[RequestType proto.Message, ResponseType proto.Message] struct {
i *info.RequestInfo
@@ -198,7 +198,7 @@ func (h *rpcHandlerImpl[RequestType, ResponseType]) claimRequest(
var affinity float32
if h.affinityFunc != nil {
affinity = h.affinityFunc(req)
affinity = h.affinityFunc(ctx, req)
if affinity < 0 {
return false, nil
}
pkg/server/stream.go
+2 -2
View File
@@ -32,7 +32,7 @@ import (
"github.com/livekit/psrpc/pkg/metadata"
)
type StreamAffinityFunc func() float32
type StreamAffinityFunc func(ctx context.Context) float32
type streamHandler[RecvType, SendType proto.Message] struct {
i *info.RequestInfo
@@ -215,7 +215,7 @@ func (h *streamHandler[RecvType, SendType]) claimRequest(
var affinity float32
if h.affinityFunc != nil {
affinity = h.affinityFunc()
affinity = h.affinityFunc(ctx)
if affinity < 0 {
return false, nil
}
protoc-gen-psrpc/generator.go
+2 -2
View File
@@ -561,12 +561,12 @@ func (t *psrpc) generateServerImplSignature(method *descriptor.MethodDescriptorP
if opts.Stream {
t.P(` `, methName, `(`, t.pkgs["psrpc"], `.ServerStream[*`, outputType, `, *`, inputType, `]) error`)
if opts.Type == options.Routing_AFFINITY {
t.P(` `, methName, `Affinity() float32`)
t.P(` `, methName, `Affinity(context.Context) float32`)
}
} else {
t.P(` `, methName, `(`, t.pkgs["context"], `.Context, *`, inputType, `) (*`, outputType, `, error)`)
if opts.Type == options.Routing_AFFINITY {
t.P(` `, methName, `Affinity(*`, inputType, `) float32`)
t.P(` `, methName, `Affinity(context.Context, *`, inputType, `) float32`)
}
}
t.P()
version/version.go
+2 -2
View File
@@ -15,6 +15,6 @@
package version
const (
Version = "v0.3.3"
PsrpcVersion_0_3_3 = true
Version = "v0.3.4"
PsrpcVersion_0_3_4 = true
)