fswriters.go 3.8 KB

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