-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
worker_calltracker.go
117 lines (90 loc) · 2.14 KB
/
worker_calltracker.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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package sectorstorage
import (
"fmt"
"io"
"github.com/filecoin-project/go-statestore"
cbg "github.com/whyrusleeping/cbor-gen"
"golang.org/x/xerrors"
"github.com/filecoin-project/lotus/extern/sector-storage/storiface"
)
type workerCallTracker struct {
st *statestore.StateStore // by CallID
}
type CallState uint64
const (
CallStarted CallState = iota
CallDone
// returned -> remove
)
type Call struct {
ID storiface.CallID
RetType ReturnType
State CallState
Result *ManyBytes // json bytes
}
func (wt *workerCallTracker) onStart(ci storiface.CallID, rt ReturnType) error {
return wt.st.Begin(ci, &Call{
ID: ci,
RetType: rt,
State: CallStarted,
})
}
func (wt *workerCallTracker) onDone(ci storiface.CallID, ret []byte) error {
st := wt.st.Get(ci)
return st.Mutate(func(cs *Call) error {
cs.State = CallDone
cs.Result = &ManyBytes{ret}
return nil
})
}
func (wt *workerCallTracker) onReturned(ci storiface.CallID) error {
st := wt.st.Get(ci)
return st.End()
}
func (wt *workerCallTracker) unfinished() ([]Call, error) {
var out []Call
return out, wt.st.List(&out)
}
// Ideally this would be a tag on the struct field telling cbor-gen to enforce higher max-len
type ManyBytes struct {
b []byte
}
const many = 100 << 20
func (t *ManyBytes) MarshalCBOR(w io.Writer) error {
if t == nil {
t = &ManyBytes{}
}
if len(t.b) > many {
return xerrors.Errorf("byte array in field t.Result was too long")
}
scratch := make([]byte, 9)
if err := cbg.WriteMajorTypeHeaderBuf(scratch, w, cbg.MajByteString, uint64(len(t.b))); err != nil {
return err
}
if _, err := w.Write(t.b[:]); err != nil {
return err
}
return nil
}
func (t *ManyBytes) UnmarshalCBOR(r io.Reader) error {
*t = ManyBytes{}
br := cbg.GetPeeker(r)
scratch := make([]byte, 9)
maj, extra, err := cbg.CborReadHeaderBuf(br, scratch)
if err != nil {
return err
}
if extra > many {
return fmt.Errorf("byte array too large (%d)", extra)
}
if maj != cbg.MajByteString {
return fmt.Errorf("expected byte array")
}
if extra > 0 {
t.b = make([]uint8, extra)
}
if _, err := io.ReadFull(br, t.b[:]); err != nil {
return err
}
return nil
}