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
|
package util
import (
"context"
"encoding/binary"
"io"
"sync"
"github.com/tonistiigi/fsutil"
)
var bufPool = sync.Pool{
New: func() interface{} {
return make([]byte, 32*1<<10)
},
}
func NewProtoStream(ctx context.Context, r io.Reader, w io.Writer) fsutil.Stream {
return &protoStream{ctx, r, w}
}
type protoStream struct {
ctx context.Context
io.Reader
io.Writer
}
func (c *protoStream) RecvMsg(m interface{}) error {
type unmarshaler interface {
Unmarshal([]byte) error
}
var h [4]byte
if _, err := io.ReadFull(c.Reader, h[:]); err != nil {
return err
}
msg := m.(unmarshaler)
length := binary.BigEndian.Uint32(h[:])
if length == 0 {
return nil
}
buf := bufPool.Get().([]byte)
if cap(buf) < int(length) {
buf = make([]byte, length)
} else {
buf = buf[:length]
}
defer bufPool.Put(buf)
if _, err := io.ReadFull(c.Reader, buf); err != nil {
return err
}
err := msg.Unmarshal(buf)
if err != nil {
return err
}
return nil
}
func (fc *protoStream) SendMsg(m interface{}) error {
type marshalerSizer interface {
MarshalTo([]byte) (int, error)
Size() int
}
msg := m.(marshalerSizer)
size := msg.Size()
b := make([]byte, msg.Size()+4)
binary.BigEndian.PutUint32(b[:4], uint32(size))
if _, err := msg.MarshalTo(b[4:]); err != nil {
return err
}
_, err := fc.Writer.Write(b)
return err
}
func (fc *protoStream) Context() context.Context {
return fc.ctx
}
|
