github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/snapshot/snapshot.go (about) 1 // snapshot manages the interactions between Nomad and Raft in order to take 2 // and restore snapshots for disaster recovery. The internal format of a 3 // snapshot is simply a tar file, as described in archive.go. 4 package snapshot 5 6 import ( 7 "compress/gzip" 8 "crypto/sha256" 9 "encoding/base64" 10 "fmt" 11 "io" 12 "io/ioutil" 13 "os" 14 15 "github.com/hashicorp/go-hclog" 16 "github.com/hashicorp/raft" 17 ) 18 19 // Snapshot is a structure that holds state about a temporary file that is used 20 // to hold a snapshot. By using an intermediate file we avoid holding everything 21 // in memory. 22 type Snapshot struct { 23 file *os.File 24 index uint64 25 checksum string 26 } 27 28 // New takes a state snapshot of the given Raft instance into a temporary file 29 // and returns an object that gives access to the file as an io.Reader. You must 30 // arrange to call Close() on the returned object or else you will leak a 31 // temporary file. 32 func New(logger hclog.Logger, r *raft.Raft) (*Snapshot, error) { 33 // Take the snapshot. 34 future := r.Snapshot() 35 if err := future.Error(); err != nil { 36 return nil, fmt.Errorf("Raft error when taking snapshot: %v", err) 37 } 38 39 // Open up the snapshot. 40 metadata, snap, err := future.Open() 41 if err != nil { 42 return nil, fmt.Errorf("failed to open snapshot: %v:", err) 43 } 44 defer func() { 45 if err := snap.Close(); err != nil { 46 logger.Error("Failed to close Raft snapshot", "error", err) 47 } 48 }() 49 50 // Make a scratch file to receive the contents so that we don't buffer 51 // everything in memory. This gets deleted in Close() since we keep it 52 // around for re-reading. 53 archive, err := ioutil.TempFile("", "snapshot") 54 if err != nil { 55 return nil, fmt.Errorf("failed to create snapshot file: %v", err) 56 } 57 58 // If anything goes wrong after this point, we will attempt to clean up 59 // the temp file. The happy path will disarm this. 60 var keep bool 61 defer func() { 62 if keep { 63 return 64 } 65 66 if err := os.Remove(archive.Name()); err != nil { 67 logger.Error("Failed to clean up temp snapshot", "error", err) 68 } 69 }() 70 71 hash := sha256.New() 72 out := io.MultiWriter(hash, archive) 73 74 // Wrap the file writer in a gzip compressor. 75 compressor := gzip.NewWriter(out) 76 77 // Write the archive. 78 if err := write(compressor, metadata, snap); err != nil { 79 return nil, fmt.Errorf("failed to write snapshot file: %v", err) 80 } 81 82 // Finish the compressed stream. 83 if err := compressor.Close(); err != nil { 84 return nil, fmt.Errorf("failed to compress snapshot file: %v", err) 85 } 86 87 // Sync the compressed file and rewind it so it's ready to be streamed 88 // out by the caller. 89 if err := archive.Sync(); err != nil { 90 return nil, fmt.Errorf("failed to sync snapshot: %v", err) 91 } 92 if _, err := archive.Seek(0, 0); err != nil { 93 return nil, fmt.Errorf("failed to rewind snapshot: %v", err) 94 } 95 96 checksum := "sha-256=" + base64.StdEncoding.EncodeToString(hash.Sum(nil)) 97 98 keep = true 99 return &Snapshot{archive, metadata.Index, checksum}, nil 100 } 101 102 // Index returns the index of the snapshot. This is safe to call on a nil 103 // snapshot, it will just return 0. 104 func (s *Snapshot) Index() uint64 { 105 if s == nil { 106 return 0 107 } 108 return s.index 109 } 110 111 func (s *Snapshot) Checksum() string { 112 if s == nil { 113 return "" 114 } 115 return s.checksum 116 } 117 118 // Read passes through to the underlying snapshot file. This is safe to call on 119 // a nil snapshot, it will just return an EOF. 120 func (s *Snapshot) Read(p []byte) (n int, err error) { 121 if s == nil { 122 return 0, io.EOF 123 } 124 return s.file.Read(p) 125 } 126 127 // Close closes the snapshot and removes any temporary storage associated with 128 // it. You must arrange to call this whenever NewSnapshot() has been called 129 // successfully. This is safe to call on a nil snapshot. 130 func (s *Snapshot) Close() error { 131 if s == nil { 132 return nil 133 } 134 135 if err := s.file.Close(); err != nil { 136 return err 137 } 138 return os.Remove(s.file.Name()) 139 } 140 141 type Discard struct { 142 io.Writer 143 } 144 145 func (dc Discard) Close() error { return nil } 146 147 // Verify takes the snapshot from the reader and verifies its contents. 148 func Verify(in io.Reader) (*raft.SnapshotMeta, error) { 149 return CopySnapshot(in, Discard{Writer: io.Discard}) 150 } 151 152 // CopySnapshot copies the snapshot content from snapshot archive to dest. 153 // It will close the destination once complete. 154 func CopySnapshot(in io.Reader, dest io.WriteCloser) (*raft.SnapshotMeta, error) { 155 defer dest.Close() 156 157 // Wrap the reader in a gzip decompressor. 158 decomp, err := gzip.NewReader(in) 159 if err != nil { 160 return nil, fmt.Errorf("failed to decompress snapshot: %v", err) 161 } 162 defer decomp.Close() 163 164 // Read the archive, throwing away the snapshot data. 165 var metadata raft.SnapshotMeta 166 if err := read(decomp, &metadata, dest); err != nil { 167 return nil, fmt.Errorf("failed to read snapshot file: %v", err) 168 } 169 170 if err := concludeGzipRead(decomp); err != nil { 171 return nil, err 172 } 173 174 return &metadata, nil 175 } 176 177 // concludeGzipRead should be invoked after you think you've consumed all of 178 // the data from the gzip stream. It will error if the stream was corrupt. 179 // 180 // The docs for gzip.Reader say: "Clients should treat data returned by Read as 181 // tentative until they receive the io.EOF marking the end of the data." 182 func concludeGzipRead(decomp *gzip.Reader) error { 183 extra, err := ioutil.ReadAll(decomp) // ReadAll consumes the EOF 184 if err != nil { 185 return err 186 } else if len(extra) != 0 { 187 return fmt.Errorf("%d unread uncompressed bytes remain", len(extra)) 188 } 189 return nil 190 } 191 192 type readWrapper struct { 193 in io.Reader 194 c int 195 } 196 197 func (r *readWrapper) Read(b []byte) (int, error) { 198 n, err := r.in.Read(b) 199 r.c += n 200 if err != nil && err != io.EOF { 201 return n, fmt.Errorf("failed to read after %v: %v", r.c, err) 202 } 203 return n, err 204 } 205 206 // Restore takes the snapshot from the reader and attempts to apply it to the 207 // given Raft instance. 208 func Restore(logger hclog.Logger, in io.Reader, r *raft.Raft) error { 209 // Wrap the reader in a gzip decompressor. 210 decomp, err := gzip.NewReader(&readWrapper{in, 0}) 211 if err != nil { 212 return fmt.Errorf("failed to decompress snapshot: %v", err) 213 } 214 defer func() { 215 if err := decomp.Close(); err != nil { 216 logger.Error("Failed to close snapshot decompressor", "error", err) 217 } 218 }() 219 220 // Make a scratch file to receive the contents of the snapshot data so 221 // we can avoid buffering in memory. 222 snap, err := ioutil.TempFile("", "snapshot") 223 if err != nil { 224 return fmt.Errorf("failed to create temp snapshot file: %v", err) 225 } 226 defer func() { 227 if err := snap.Close(); err != nil { 228 logger.Error("Failed to close temp snapshot", "error", err) 229 } 230 if err := os.Remove(snap.Name()); err != nil { 231 logger.Error("Failed to clean up temp snapshot", "error", err) 232 } 233 }() 234 235 // Read the archive. 236 var metadata raft.SnapshotMeta 237 if err := read(decomp, &metadata, snap); err != nil { 238 return fmt.Errorf("failed to read snapshot file: %v", err) 239 } 240 241 if err := concludeGzipRead(decomp); err != nil { 242 return err 243 } 244 245 // Sync and rewind the file so it's ready to be read again. 246 if err := snap.Sync(); err != nil { 247 return fmt.Errorf("failed to sync temp snapshot: %v", err) 248 } 249 if _, err := snap.Seek(0, 0); err != nil { 250 return fmt.Errorf("failed to rewind temp snapshot: %v", err) 251 } 252 253 // Feed the snapshot into Raft. 254 if err := r.Restore(&metadata, snap, 0); err != nil { 255 return fmt.Errorf("Raft error when restoring snapshot: %v", err) 256 } 257 258 return nil 259 }