// handleSnapshotRequest reads the request from the conn and dispatches it. This // will be called from a goroutine after an incoming stream is determined to be // a snapshot request. func (s *Server) handleSnapshotRequest(conn net.Conn) error { var args structs.SnapshotRequest dec := codec.NewDecoder(conn, &codec.MsgpackHandle{}) if err := dec.Decode(&args); err != nil { return fmt.Errorf("failed to decode request: %v", err) } var reply structs.SnapshotResponse snap, err := s.dispatchSnapshotRequest(&args, conn, &reply) if err != nil { reply.Error = err.Error() goto RESPOND } defer func() { if err := snap.Close(); err != nil { s.logger.Printf("[ERR] consul: Failed to close snapshot: %v", err) } }() RESPOND: enc := codec.NewEncoder(conn, &codec.MsgpackHandle{}) if err := enc.Encode(&reply); err != nil { return fmt.Errorf("failed to encode response: %v", err) } if snap != nil { if _, err := io.Copy(conn, snap); err != nil { return fmt.Errorf("failed to stream snapshot: %v", err) } } return nil }
// dispatchSnapshotRequest takes an incoming request structure with possibly some // streaming data (for a restore) and returns possibly some streaming data (for // a snapshot save). We can't use the normal RPC mechanism in a streaming manner // like this, so we have to dispatch these by hand. func (s *Server) dispatchSnapshotRequest(args *structs.SnapshotRequest, in io.Reader, reply *structs.SnapshotResponse) (io.ReadCloser, error) { // Perform DC forwarding. if dc := args.Datacenter; dc != s.config.Datacenter { server, ok := s.getRemoteServer(dc) if !ok { return nil, structs.ErrNoDCPath } return SnapshotRPC(s.connPool, dc, server.Addr, args, in, reply) } // Perform leader forwarding if required. if !args.AllowStale { if isLeader, server := s.getLeader(); !isLeader { if server == nil { return nil, structs.ErrNoLeader } return SnapshotRPC(s.connPool, args.Datacenter, server.Addr, args, in, reply) } } // Snapshots don't work in dev mode because we need Raft's snapshots to // be readable. Better to present a clear error than one from deep down // in the Raft snapshot store. if s.config.DevMode { return nil, errors.New(noSnapshotsInDevMode) } // Verify token is allowed to operate on snapshots. There's only a // single ACL sense here (not read and write) since reading gets you // all the ACLs and you could escalate from there. if acl, err := s.resolveToken(args.Token); err != nil { return nil, err } else if acl != nil && !acl.Snapshot() { return nil, permissionDeniedErr } // Dispatch the operation. switch args.Op { case structs.SnapshotSave: if !args.AllowStale { if err := s.consistentRead(); err != nil { return nil, err } } // Set the metadata here before we do anything; this should always be // pessimistic if we get more data while the snapshot is being taken. s.setQueryMeta(&reply.QueryMeta) // Take the snapshot and capture the index. snap, err := snapshot.New(s.logger, s.raft) reply.Index = snap.Index() return snap, err case structs.SnapshotRestore: if args.AllowStale { return nil, fmt.Errorf("stale not allowed for restore") } // Restore the snapshot. if err := snapshot.Restore(s.logger, in, s.raft); err != nil { return nil, err } // Run a barrier so we are sure that our FSM is caught up with // any snapshot restore details (it's also part of Raft's restore // process but we don't want to depend on that detail for this to // be correct). Once that works, we can redo the leader actions // so our leader-maintained state will be up to date. barrier := s.raft.Barrier(0) if err := barrier.Error(); err != nil { return nil, err } if err := s.revokeLeadership(); err != nil { return nil, err } if err := s.establishLeadership(); err != nil { return nil, err } // Give the caller back an empty reader since there's nothing to // stream back. return ioutil.NopCloser(bytes.NewReader([]byte(""))), nil default: return nil, fmt.Errorf("unrecognized snapshot op %q", args.Op) } }