You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

128 lines
4.4 KiB

  1. package common
  2. import (
  3. fmt "fmt"
  4. "io"
  5. "os"
  6. "path/filepath"
  7. "strconv"
  8. "strings"
  9. "sync"
  10. "time"
  11. )
  12. const (
  13. atomicWriteFilePrefix = "write-file-atomic-"
  14. // Maximum number of atomic write file conflicts before we start reseeding
  15. // (reduced from golang's default 10 due to using an increased randomness space)
  16. atomicWriteFileMaxNumConflicts = 5
  17. // Maximum number of attempts to make at writing the write file before giving up
  18. // (reduced from golang's default 10000 due to using an increased randomness space)
  19. atomicWriteFileMaxNumWriteAttempts = 1000
  20. // LCG constants from Donald Knuth MMIX
  21. // This LCG's has a period equal to 2**64
  22. lcgA = 6364136223846793005
  23. lcgC = 1442695040888963407
  24. // Create in case it doesn't exist and force kernel
  25. // flush, which still leaves the potential of lingering disk cache.
  26. // Never overwrites files
  27. atomicWriteFileFlag = os.O_WRONLY | os.O_CREATE | os.O_SYNC | os.O_TRUNC | os.O_EXCL
  28. )
  29. var (
  30. atomicWriteFileRand uint64
  31. atomicWriteFileRandMu sync.Mutex
  32. )
  33. func writeFileRandReseed() uint64 {
  34. // Scale the PID, to minimize the chance that two processes seeded at similar times
  35. // don't get the same seed. Note that PID typically ranges in [0, 2**15), but can be
  36. // up to 2**22 under certain configurations. We left bit-shift the PID by 20, so that
  37. // a PID difference of one corresponds to a time difference of 2048 seconds.
  38. // The important thing here is that now for a seed conflict, they would both have to be on
  39. // the correct nanosecond offset, and second-based offset, which is much less likely than
  40. // just a conflict with the correct nanosecond offset.
  41. return uint64(time.Now().UnixNano() + int64(os.Getpid()<<20))
  42. }
  43. // Use a fast thread safe LCG for atomic write file names.
  44. // Returns a string corresponding to a 64 bit int.
  45. // If it was a negative int, the leading number is a 0.
  46. func randWriteFileSuffix() string {
  47. atomicWriteFileRandMu.Lock()
  48. r := atomicWriteFileRand
  49. if r == 0 {
  50. r = writeFileRandReseed()
  51. }
  52. // Update randomness according to lcg
  53. r = r*lcgA + lcgC
  54. atomicWriteFileRand = r
  55. atomicWriteFileRandMu.Unlock()
  56. // Can have a negative name, replace this in the following
  57. suffix := strconv.Itoa(int(r))
  58. if string(suffix[0]) == "-" {
  59. // Replace first "-" with "0". This is purely for UI clarity,
  60. // as otherwhise there would be two `-` in a row.
  61. suffix = strings.Replace(suffix, "-", "0", 1)
  62. }
  63. return suffix
  64. }
  65. // WriteFileAtomic creates a temporary file with data and provided perm and
  66. // swaps it atomically with filename if successful.
  67. func WriteFileAtomic(filename string, data []byte, perm os.FileMode) (err error) {
  68. // This implementation is inspired by the golang stdlibs method of creating
  69. // tempfiles. Notable differences are that we use different flags, a 64 bit LCG
  70. // and handle negatives differently.
  71. // The core reason we can't use golang's TempFile is that we must write
  72. // to the file synchronously, as we need this to persist to disk.
  73. // We also open it in write-only mode, to avoid concerns that arise with read.
  74. var (
  75. dir = filepath.Dir(filename)
  76. f *os.File
  77. )
  78. nconflict := 0
  79. // Limit the number of attempts to create a file. Something is seriously
  80. // wrong if it didn't get created after 1000 attempts, and we don't want
  81. // an infinite loop
  82. i := 0
  83. for ; i < atomicWriteFileMaxNumWriteAttempts; i++ {
  84. name := filepath.Join(dir, atomicWriteFilePrefix+randWriteFileSuffix())
  85. f, err = os.OpenFile(name, atomicWriteFileFlag, perm)
  86. // If the file already exists, try a new file
  87. if os.IsExist(err) {
  88. // If the files exists too many times, start reseeding as we've
  89. // likely hit another instances seed.
  90. if nconflict++; nconflict > atomicWriteFileMaxNumConflicts {
  91. atomicWriteFileRandMu.Lock()
  92. atomicWriteFileRand = writeFileRandReseed()
  93. atomicWriteFileRandMu.Unlock()
  94. }
  95. continue
  96. } else if err != nil {
  97. return err
  98. }
  99. break
  100. }
  101. if i == atomicWriteFileMaxNumWriteAttempts {
  102. return fmt.Errorf("could not create atomic write file after %d attempts", i)
  103. }
  104. // Clean up in any case. Defer stacking order is last-in-first-out.
  105. defer os.Remove(f.Name())
  106. defer f.Close()
  107. if n, err := f.Write(data); err != nil {
  108. return err
  109. } else if n < len(data) {
  110. return io.ErrShortWrite
  111. }
  112. // Close the file before renaming it, otherwise it will cause "The process
  113. // cannot access the file because it is being used by another process." on windows.
  114. f.Close()
  115. return os.Rename(f.Name(), filename)
  116. }