github.com/criteo-forks/consul@v1.4.5-criteonogrpc/snapshot/snapshot.go (about) 1 // snapshot manages the interactions between Consul 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 "fmt" 9 "io" 10 "io/ioutil" 11 "log" 12 "os" 13 14 "github.com/hashicorp/raft" 15 ) 16 17 // Snapshot is a structure that holds state about a temporary file that is used 18 // to hold a snapshot. By using an intermediate file we avoid holding everything 19 // in memory. 20 type Snapshot struct { 21 file *os.File 22 index uint64 23 } 24 25 // New takes a state snapshot of the given Raft instance into a temporary file 26 // and returns an object that gives access to the file as an io.Reader. You must 27 // arrange to call Close() on the returned object or else you will leak a 28 // temporary file. 29 func New(logger *log.Logger, r *raft.Raft) (*Snapshot, error) { 30 // Take the snapshot. 31 future := r.Snapshot() 32 if err := future.Error(); err != nil { 33 return nil, fmt.Errorf("Raft error when taking snapshot: %v", err) 34 } 35 36 // Open up the snapshot. 37 metadata, snap, err := future.Open() 38 if err != nil { 39 return nil, fmt.Errorf("failed to open snapshot: %v:", err) 40 } 41 defer func() { 42 if err := snap.Close(); err != nil { 43 logger.Printf("[ERR] snapshot: Failed to close Raft snapshot: %v", err) 44 } 45 }() 46 47 // Make a scratch file to receive the contents so that we don't buffer 48 // everything in memory. This gets deleted in Close() since we keep it 49 // around for re-reading. 50 archive, err := ioutil.TempFile("", "snapshot") 51 if err != nil { 52 return nil, fmt.Errorf("failed to create snapshot file: %v", err) 53 } 54 55 // If anything goes wrong after this point, we will attempt to clean up 56 // the temp file. The happy path will disarm this. 57 var keep bool 58 defer func() { 59 if keep { 60 return 61 } 62 63 if err := os.Remove(archive.Name()); err != nil { 64 logger.Printf("[ERR] snapshot: Failed to clean up temp snapshot: %v", err) 65 } 66 }() 67 68 // Wrap the file writer in a gzip compressor. 69 compressor := gzip.NewWriter(archive) 70 71 // Write the archive. 72 if err := write(compressor, metadata, snap); err != nil { 73 return nil, fmt.Errorf("failed to write snapshot file: %v", err) 74 } 75 76 // Finish the compressed stream. 77 if err := compressor.Close(); err != nil { 78 return nil, fmt.Errorf("failed to compress snapshot file: %v", err) 79 } 80 81 // Sync the compressed file and rewind it so it's ready to be streamed 82 // out by the caller. 83 if err := archive.Sync(); err != nil { 84 return nil, fmt.Errorf("failed to sync snapshot: %v", err) 85 } 86 if _, err := archive.Seek(0, 0); err != nil { 87 return nil, fmt.Errorf("failed to rewind snapshot: %v", err) 88 } 89 90 keep = true 91 return &Snapshot{archive, metadata.Index}, nil 92 } 93 94 // Index returns the index of the snapshot. This is safe to call on a nil 95 // snapshot, it will just return 0. 96 func (s *Snapshot) Index() uint64 { 97 if s == nil { 98 return 0 99 } 100 return s.index 101 } 102 103 // Read passes through to the underlying snapshot file. This is safe to call on 104 // a nil snapshot, it will just return an EOF. 105 func (s *Snapshot) Read(p []byte) (n int, err error) { 106 if s == nil { 107 return 0, io.EOF 108 } 109 return s.file.Read(p) 110 } 111 112 // Close closes the snapshot and removes any temporary storage associated with 113 // it. You must arrange to call this whenever NewSnapshot() has been called 114 // successfully. This is safe to call on a nil snapshot. 115 func (s *Snapshot) Close() error { 116 if s == nil { 117 return nil 118 } 119 120 if err := s.file.Close(); err != nil { 121 return err 122 } 123 return os.Remove(s.file.Name()) 124 } 125 126 // Verify takes the snapshot from the reader and verifies its contents. 127 func Verify(in io.Reader) (*raft.SnapshotMeta, error) { 128 // Wrap the reader in a gzip decompressor. 129 decomp, err := gzip.NewReader(in) 130 if err != nil { 131 return nil, fmt.Errorf("failed to decompress snapshot: %v", err) 132 } 133 defer decomp.Close() 134 135 // Read the archive, throwing away the snapshot data. 136 var metadata raft.SnapshotMeta 137 if err := read(decomp, &metadata, ioutil.Discard); err != nil { 138 return nil, fmt.Errorf("failed to read snapshot file: %v", err) 139 } 140 return &metadata, nil 141 } 142 143 // Restore takes the snapshot from the reader and attempts to apply it to the 144 // given Raft instance. 145 func Restore(logger *log.Logger, in io.Reader, r *raft.Raft) error { 146 // Wrap the reader in a gzip decompressor. 147 decomp, err := gzip.NewReader(in) 148 if err != nil { 149 return fmt.Errorf("failed to decompress snapshot: %v", err) 150 } 151 defer func() { 152 if err := decomp.Close(); err != nil { 153 logger.Printf("[ERR] snapshot: Failed to close snapshot decompressor: %v", err) 154 } 155 }() 156 157 // Make a scratch file to receive the contents of the snapshot data so 158 // we can avoid buffering in memory. 159 snap, err := ioutil.TempFile("", "snapshot") 160 if err != nil { 161 return fmt.Errorf("failed to create temp snapshot file: %v", err) 162 } 163 defer func() { 164 if err := snap.Close(); err != nil { 165 logger.Printf("[ERR] snapshot: Failed to close temp snapshot: %v", err) 166 } 167 if err := os.Remove(snap.Name()); err != nil { 168 logger.Printf("[ERR] snapshot: Failed to clean up temp snapshot: %v", err) 169 } 170 }() 171 172 // Read the archive. 173 var metadata raft.SnapshotMeta 174 if err := read(decomp, &metadata, snap); err != nil { 175 return fmt.Errorf("failed to read snapshot file: %v", err) 176 } 177 178 // Sync and rewind the file so it's ready to be read again. 179 if err := snap.Sync(); err != nil { 180 return fmt.Errorf("failed to sync temp snapshot: %v", err) 181 } 182 if _, err := snap.Seek(0, 0); err != nil { 183 return fmt.Errorf("failed to rewind temp snapshot: %v", err) 184 } 185 186 // Feed the snapshot into Raft. 187 if err := r.Restore(&metadata, snap, 0); err != nil { 188 return fmt.Errorf("Raft error when restoring snapshot: %v", err) 189 } 190 191 return nil 192 }