-
Notifications
You must be signed in to change notification settings - Fork 0
/
stub.go
68 lines (54 loc) · 1.31 KB
/
stub.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
package call
import (
"context"
"github.com/kanengo/akasar/runtime/codegen"
"github.com/kanengo/akasar/runtime/pool"
)
type stub struct {
conn Connection
methods []stubMethod
injectRetries int
}
func (s *stub) Invoke(ctx context.Context, method int, args []byte, shardKey uint64) (result []byte, err error) {
m := s.methods[method]
opts := CallOptions{
Retry: m.retry,
ShardKey: shardKey,
}
n := 1
if m.retry {
n += s.injectRetries // fake retries for testing
}
defer func() {
_ = pool.FreePowerOfTwoSizeBytes(args)
}()
for i := 0; i < n; i++ {
result, err = s.conn.Call(ctx, m.key, args, opts)
}
return
}
type stubMethod struct {
key MethodKey
retry bool
}
var _ codegen.Stub = &stub{}
func NewStub(name string, reg *codegen.Registration, conn Connection, injectRetries int) codegen.Stub {
return &stub{
conn: conn,
methods: makeStubMethods(name, reg),
injectRetries: injectRetries,
}
}
func makeStubMethods(fullName string, reg *codegen.Registration) []stubMethod {
n := reg.Iface.NumMethod()
methods := make([]stubMethod, n)
for i := 0; i < n; i++ {
mName := reg.Iface.Method(i).Name
methods[i].key = MakeMethodKey(fullName, mName)
methods[i].retry = true
}
for _, m := range reg.NoRetry {
methods[m].retry = true
}
return methods
}