/
rpc.go
101 lines (80 loc) · 2.03 KB
/
rpc.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
// Package rpc implements some of the lower-level functionality required to
// communicate with the namenode and datanodes.
package rpc
import (
"encoding/binary"
"errors"
"io"
"math/rand"
"time"
"google.golang.org/protobuf/proto"
)
var errInvalidResponse = errors.New("invalid response from namenode")
// Used for client ID generation, below.
const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
func newClientID() []byte {
id := make([]byte, 16)
rand.Seed(time.Now().UTC().UnixNano())
for i := range id {
id[i] = chars[rand.Intn(len(chars))]
}
return id
}
func makePrefixedMessage(msg proto.Message) ([]byte, error) {
msgBytes, err := proto.Marshal(msg)
if err != nil {
return nil, err
}
lengthBytes := make([]byte, 10)
n := binary.PutUvarint(lengthBytes, uint64(len(msgBytes)))
return append(lengthBytes[:n], msgBytes...), nil
}
func makeRPCPacket(msgs ...proto.Message) ([]byte, error) {
packet := make([]byte, 4, 128)
length := 0
for _, msg := range msgs {
b, err := makePrefixedMessage(msg)
if err != nil {
return nil, err
}
packet = append(packet, b...)
length += len(b)
}
binary.BigEndian.PutUint32(packet, uint32(length))
return packet, nil
}
func readRPCPacket(r io.Reader, msgs ...proto.Message) error {
var packetLength uint32
err := binary.Read(r, binary.BigEndian, &packetLength)
if err != nil {
return err
}
packet := make([]byte, packetLength)
_, err = io.ReadFull(r, packet)
if err != nil {
return err
}
for _, msg := range msgs {
// HDFS doesn't send all the response messages all the time (for example, if
// the RpcResponseHeaderProto contains an error).
if len(packet) == 0 {
return nil
}
msgLength, n := binary.Uvarint(packet)
if n <= 0 || msgLength > uint64(len(packet)) {
return errInvalidResponse
}
packet = packet[n:]
if msgLength != 0 {
err = proto.Unmarshal(packet[:msgLength], msg)
if err != nil {
return err
}
packet = packet[msgLength:]
}
}
if len(packet) > 0 {
return errInvalidResponse
}
return nil
}