1 // 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 15 package snap 16 17 import ( 18 "io" 19 20 "go.etcd.io/etcd/pkg/v3/ioutil" 21 "go.etcd.io/etcd/raft/v3/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. 32 type Message struct { 33 raftpb.Message 34 ReadCloser io.ReadCloser 35 TotalSize int64 36 closeC chan bool 37 } 38 39 func 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. 51 func (m Message) CloseNotify() <-chan bool { 52 return m.closeC 53 } 54 55 func (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 } 65