David K. Bainbridge | 215e024 | 2017-09-05 23:18:24 -0700 | [diff] [blame] | 1 | package ioutils |
| 2 | |
| 3 | import ( |
| 4 | "io" |
| 5 | "io/ioutil" |
| 6 | "os" |
| 7 | "path/filepath" |
| 8 | ) |
| 9 | |
| 10 | // NewAtomicFileWriter returns WriteCloser so that writing to it writes to a |
| 11 | // temporary file and closing it atomically changes the temporary file to |
| 12 | // destination path. Writing and closing concurrently is not allowed. |
| 13 | func NewAtomicFileWriter(filename string, perm os.FileMode) (io.WriteCloser, error) { |
| 14 | f, err := ioutil.TempFile(filepath.Dir(filename), ".tmp-"+filepath.Base(filename)) |
| 15 | if err != nil { |
| 16 | return nil, err |
| 17 | } |
| 18 | |
| 19 | abspath, err := filepath.Abs(filename) |
| 20 | if err != nil { |
| 21 | return nil, err |
| 22 | } |
| 23 | return &atomicFileWriter{ |
| 24 | f: f, |
| 25 | fn: abspath, |
| 26 | perm: perm, |
| 27 | }, nil |
| 28 | } |
| 29 | |
| 30 | // AtomicWriteFile atomically writes data to a file named by filename. |
| 31 | func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error { |
| 32 | f, err := NewAtomicFileWriter(filename, perm) |
| 33 | if err != nil { |
| 34 | return err |
| 35 | } |
| 36 | n, err := f.Write(data) |
| 37 | if err == nil && n < len(data) { |
| 38 | err = io.ErrShortWrite |
| 39 | f.(*atomicFileWriter).writeErr = err |
| 40 | } |
| 41 | if err1 := f.Close(); err == nil { |
| 42 | err = err1 |
| 43 | } |
| 44 | return err |
| 45 | } |
| 46 | |
| 47 | type atomicFileWriter struct { |
| 48 | f *os.File |
| 49 | fn string |
| 50 | writeErr error |
| 51 | perm os.FileMode |
| 52 | } |
| 53 | |
| 54 | func (w *atomicFileWriter) Write(dt []byte) (int, error) { |
| 55 | n, err := w.f.Write(dt) |
| 56 | if err != nil { |
| 57 | w.writeErr = err |
| 58 | } |
| 59 | return n, err |
| 60 | } |
| 61 | |
| 62 | func (w *atomicFileWriter) Close() (retErr error) { |
| 63 | defer func() { |
| 64 | if retErr != nil || w.writeErr != nil { |
| 65 | os.Remove(w.f.Name()) |
| 66 | } |
| 67 | }() |
| 68 | if err := w.f.Sync(); err != nil { |
| 69 | w.f.Close() |
| 70 | return err |
| 71 | } |
| 72 | if err := w.f.Close(); err != nil { |
| 73 | return err |
| 74 | } |
| 75 | if err := os.Chmod(w.f.Name(), w.perm); err != nil { |
| 76 | return err |
| 77 | } |
| 78 | if w.writeErr == nil { |
| 79 | return os.Rename(w.f.Name(), w.fn) |
| 80 | } |
| 81 | return nil |
| 82 | } |
| 83 | |
| 84 | // AtomicWriteSet is used to atomically write a set |
| 85 | // of files and ensure they are visible at the same time. |
| 86 | // Must be committed to a new directory. |
| 87 | type AtomicWriteSet struct { |
| 88 | root string |
| 89 | } |
| 90 | |
| 91 | // NewAtomicWriteSet creates a new atomic write set to |
| 92 | // atomically create a set of files. The given directory |
| 93 | // is used as the base directory for storing files before |
| 94 | // commit. If no temporary directory is given the system |
| 95 | // default is used. |
| 96 | func NewAtomicWriteSet(tmpDir string) (*AtomicWriteSet, error) { |
| 97 | td, err := ioutil.TempDir(tmpDir, "write-set-") |
| 98 | if err != nil { |
| 99 | return nil, err |
| 100 | } |
| 101 | |
| 102 | return &AtomicWriteSet{ |
| 103 | root: td, |
| 104 | }, nil |
| 105 | } |
| 106 | |
| 107 | // WriteFile writes a file to the set, guaranteeing the file |
| 108 | // has been synced. |
| 109 | func (ws *AtomicWriteSet) WriteFile(filename string, data []byte, perm os.FileMode) error { |
| 110 | f, err := ws.FileWriter(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm) |
| 111 | if err != nil { |
| 112 | return err |
| 113 | } |
| 114 | n, err := f.Write(data) |
| 115 | if err == nil && n < len(data) { |
| 116 | err = io.ErrShortWrite |
| 117 | } |
| 118 | if err1 := f.Close(); err == nil { |
| 119 | err = err1 |
| 120 | } |
| 121 | return err |
| 122 | } |
| 123 | |
| 124 | type syncFileCloser struct { |
| 125 | *os.File |
| 126 | } |
| 127 | |
| 128 | func (w syncFileCloser) Close() error { |
| 129 | err := w.File.Sync() |
| 130 | if err1 := w.File.Close(); err == nil { |
| 131 | err = err1 |
| 132 | } |
| 133 | return err |
| 134 | } |
| 135 | |
| 136 | // FileWriter opens a file writer inside the set. The file |
| 137 | // should be synced and closed before calling commit. |
| 138 | func (ws *AtomicWriteSet) FileWriter(name string, flag int, perm os.FileMode) (io.WriteCloser, error) { |
| 139 | f, err := os.OpenFile(filepath.Join(ws.root, name), flag, perm) |
| 140 | if err != nil { |
| 141 | return nil, err |
| 142 | } |
| 143 | return syncFileCloser{f}, nil |
| 144 | } |
| 145 | |
| 146 | // Cancel cancels the set and removes all temporary data |
| 147 | // created in the set. |
| 148 | func (ws *AtomicWriteSet) Cancel() error { |
| 149 | return os.RemoveAll(ws.root) |
| 150 | } |
| 151 | |
| 152 | // Commit moves all created files to the target directory. The |
| 153 | // target directory must not exist and the parent of the target |
| 154 | // directory must exist. |
| 155 | func (ws *AtomicWriteSet) Commit(target string) error { |
| 156 | return os.Rename(ws.root, target) |
| 157 | } |
| 158 | |
| 159 | // String returns the location the set is writing to. |
| 160 | func (ws *AtomicWriteSet) String() string { |
| 161 | return ws.root |
| 162 | } |