blob: c1151e27eb13af71129217d8ad99b0bf098f7485 [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// Copyright 2015 The etcd Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package snap
16
17import (
18 "io"
19
20 "go.etcd.io/etcd/pkg/ioutil"
21 "go.etcd.io/etcd/raft/raftpb"
22)
23
24// Message is a struct that contains a raft Message and a ReadCloser. The type
25// of raft message MUST be MsgSnap, which contains the raft meta-data and an
26// additional data []byte field that contains the snapshot of the actual state
27// machine.
28// Message contains the ReadCloser field for handling large snapshot. This avoid
29// copying the entire snapshot into a byte array, which consumes a lot of memory.
30//
31// User of Message should close the Message after sending it.
32type Message struct {
33 raftpb.Message
34 ReadCloser io.ReadCloser
35 TotalSize int64
36 closeC chan bool
37}
38
39func NewMessage(rs raftpb.Message, rc io.ReadCloser, rcSize int64) *Message {
40 return &Message{
41 Message: rs,
42 ReadCloser: ioutil.NewExactReadCloser(rc, rcSize),
43 TotalSize: int64(rs.Size()) + rcSize,
44 closeC: make(chan bool, 1),
45 }
46}
47
48// CloseNotify returns a channel that receives a single value
49// when the message sent is finished. true indicates the sent
50// is successful.
51func (m Message) CloseNotify() <-chan bool {
52 return m.closeC
53}
54
55func (m Message) CloseWithError(err error) {
56 if cerr := m.ReadCloser.Close(); cerr != nil {
57 err = cerr
58 }
59 if err == nil {
60 m.closeC <- true
61 } else {
62 m.closeC <- false
63 }
64}