github.com/aergoio/aergo@v1.3.1/p2p/raftsupport/snapshotreceiver.go (about)

     1  /*
     2   * @file
     3   * @copyright defined in aergo/LICENSE.txt
     4   */
     5  
     6  package raftsupport
     7  
     8  import (
     9  	"context"
    10  	"encoding/binary"
    11  	"github.com/aergoio/aergo-lib/log"
    12  	"github.com/aergoio/aergo/consensus"
    13  	"github.com/aergoio/aergo/p2p/p2pcommon"
    14  	"github.com/aergoio/aergo/p2p/p2putil"
    15  	"github.com/aergoio/aergo/types"
    16  	rtypes "github.com/aergoio/etcd/pkg/types"
    17  	"github.com/aergoio/etcd/raft/raftpb"
    18  	"github.com/golang/protobuf/proto"
    19  	"io"
    20  )
    21  
    22  const (
    23  	SnapRespHeaderLength = 4
    24  )
    25  // TODO consider the scope of type
    26  type snapshotReceiver struct {
    27  	logger *log.Logger
    28  	pm     p2pcommon.PeerManager
    29  	rAcc   consensus.AergoRaftAccessor
    30  	peer   p2pcommon.RemotePeer
    31  	rwc    io.ReadWriteCloser
    32  }
    33  
    34  func newSnapshotReceiver(logger *log.Logger, pm p2pcommon.PeerManager, rAcc consensus.AergoRaftAccessor, peer p2pcommon.RemotePeer, sender io.ReadWriteCloser) *snapshotReceiver {
    35  	return &snapshotReceiver{logger: logger, pm: pm, rAcc: rAcc, peer: peer, rwc: sender}
    36  }
    37  
    38  
    39  func (s *snapshotReceiver) Receive() {
    40  	resp := &types.SnapshotResponse{Status:types.ResultStatus_OK}
    41  	defer s.sendResp(s.rwc, resp)
    42  
    43  	dec := &RaftMsgDecoder{r: s.rwc}
    44  	// let snapshots be very large since they can exceed 512MB for large installations
    45  	m, err := dec.DecodeLimit(uint64(1 << 63))
    46  	from := rtypes.ID(m.From).String()
    47  	if err != nil {
    48  		s.logger.Error().Str(p2putil.LogPeerName, s.peer.Name()).Err(err).Msg("failed to decode raft message")
    49  		resp.Status = types.ResultStatus_INVALID_ARGUMENT
    50  		resp.Message = "malformed message"
    51  		// TODO return error
    52  		//recvFailures.WithLabelValues(rwc.RemoteAddr).Inc()
    53  		//snapshotReceiveFailures.WithLabelValues(from).Inc()
    54  		return
    55  	}
    56  
    57  	//receivedBytes.WithLabelValues(from).Add(float64(m.Size()))
    58  
    59  	if m.Type != raftpb.MsgSnap {
    60  		s.logger.Error().Str("type", m.Type.String()).Msg("unexpected raft message type on snapshot path")
    61  		resp.Status = types.ResultStatus_INVALID_ARGUMENT
    62  		resp.Message = "invalid message type"
    63  
    64  		//http.Error(w, "wrong raft message type", http.StatusBadRequest)
    65  		//snapshotReceiveFailures.WithLabelValues(from).Inc()
    66  		return
    67  	}
    68  
    69  	s.logger.Info().Uint64("index", m.Snapshot.Metadata.Index).Str("from", from).Msg("receiving database snapshot")
    70  	// save incoming database snapshot.
    71  	_, err = s.rAcc.SaveFromRemote(s.rwc, m.Snapshot.Metadata.Index, m)
    72  	if err != nil {
    73  		s.logger.Error().Err(err).Msg("failed to save KV snapshot")
    74  		resp.Status = types.ResultStatus_INTERNAL
    75  
    76  		//http.Error(w, msg, http.StatusInternalServerError)
    77  		//snapshotReceiveFailures.WithLabelValues(from).Inc()
    78  		return
    79  	}
    80  	//receivedBytes.WithLabelValues(from).Add(float64(n))
    81  	s.logger.Info().Str(p2putil.LogPeerName, s.peer.Name()).Uint64("index", m.Snapshot.Metadata.Index).Str("from", from).Msg("received and saved database snapshot successfully")
    82  
    83  	if err := s.rAcc.Process(context.TODO(),s.peer.ID(), m); err != nil {
    84  		switch v := err.(type) {
    85  		// Process may return codeError error when doing some
    86  		// additional checks before calling raft.Node.Step.
    87  		case codeError:
    88  			// TODO get resp
    89  			resp.Status =v.Status()
    90  			resp.Message = v.Message()
    91  		default:
    92  			s.logger.Warn().Err(err).Msg("failed to process raft message")
    93  			resp.Status = types.ResultStatus_UNKNOWN
    94  			//http.Error(w, msg, http.StatusInternalServerError)
    95  			//snapshotReceiveFailures.WithLabelValues(from).Inc()
    96  		}
    97  		return
    98  	}
    99  	// Write StatusNoContent header after the message has been processed by
   100  	// raft, which facilitates the client to report MsgSnap status.
   101  	//w.WriteHeader(http.StatusNoContent)
   102  
   103  	//snapshotReceive.WithLabelValues(from).Inc()
   104  	//snapshotReceiveSeconds.WithLabelValues(from).Observe(time.Since(start).Seconds())
   105  }
   106  
   107  func (s *snapshotReceiver) sendResp(w io.Writer, resp *types.SnapshotResponse) {
   108  	b, err := proto.Marshal(resp)
   109  	if err == nil {
   110  		bytebuf := make([]byte, SnapRespHeaderLength)
   111  		binary.BigEndian.PutUint32(bytebuf, uint32(len(b)))
   112  		w.Write(bytebuf)
   113  		w.Write(b)
   114  	} else {
   115  		s.logger.Info().Err(err).Msg("Failed to write snapshot response")
   116  	}
   117  }
   118  
   119  type codeError interface {
   120  	Status() types.ResultStatus
   121  	Message() string
   122  }