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
128 lines
4.4 KiB
package tempfile
|
|
|
|
import (
|
|
fmt "fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
atomicWriteFilePrefix = "write-file-atomic-"
|
|
// Maximum number of atomic write file conflicts before we start reseeding
|
|
// (reduced from golang's default 10 due to using an increased randomness space)
|
|
atomicWriteFileMaxNumConflicts = 5
|
|
// Maximum number of attempts to make at writing the write file before giving up
|
|
// (reduced from golang's default 10000 due to using an increased randomness space)
|
|
atomicWriteFileMaxNumWriteAttempts = 1000
|
|
// LCG constants from Donald Knuth MMIX
|
|
// This LCG's has a period equal to 2**64
|
|
lcgA = 6364136223846793005
|
|
lcgC = 1442695040888963407
|
|
// Create in case it doesn't exist and force kernel
|
|
// flush, which still leaves the potential of lingering disk cache.
|
|
// Never overwrites files
|
|
atomicWriteFileFlag = os.O_WRONLY | os.O_CREATE | os.O_SYNC | os.O_TRUNC | os.O_EXCL
|
|
)
|
|
|
|
var (
|
|
atomicWriteFileRand uint64
|
|
atomicWriteFileRandMu sync.Mutex
|
|
)
|
|
|
|
func writeFileRandReseed() uint64 {
|
|
// Scale the PID, to minimize the chance that two processes seeded at similar times
|
|
// don't get the same seed. Note that PID typically ranges in [0, 2**15), but can be
|
|
// up to 2**22 under certain configurations. We left bit-shift the PID by 20, so that
|
|
// a PID difference of one corresponds to a time difference of 2048 seconds.
|
|
// The important thing here is that now for a seed conflict, they would both have to be on
|
|
// the correct nanosecond offset, and second-based offset, which is much less likely than
|
|
// just a conflict with the correct nanosecond offset.
|
|
return uint64(time.Now().UnixNano() + int64(os.Getpid()<<20))
|
|
}
|
|
|
|
// Use a fast thread safe LCG for atomic write file names.
|
|
// Returns a string corresponding to a 64 bit int.
|
|
// If it was a negative int, the leading number is a 0.
|
|
func randWriteFileSuffix() string {
|
|
atomicWriteFileRandMu.Lock()
|
|
r := atomicWriteFileRand
|
|
if r == 0 {
|
|
r = writeFileRandReseed()
|
|
}
|
|
|
|
// Update randomness according to lcg
|
|
r = r*lcgA + lcgC
|
|
|
|
atomicWriteFileRand = r
|
|
atomicWriteFileRandMu.Unlock()
|
|
// Can have a negative name, replace this in the following
|
|
suffix := strconv.Itoa(int(r))
|
|
if string(suffix[0]) == "-" {
|
|
// Replace first "-" with "0". This is purely for UI clarity,
|
|
// as otherwhise there would be two `-` in a row.
|
|
suffix = strings.Replace(suffix, "-", "0", 1)
|
|
}
|
|
return suffix
|
|
}
|
|
|
|
// WriteFileAtomic creates a temporary file with data and provided perm and
|
|
// swaps it atomically with filename if successful.
|
|
func WriteFileAtomic(filename string, data []byte, perm os.FileMode) (err error) {
|
|
// This implementation is inspired by the golang stdlibs method of creating
|
|
// tempfiles. Notable differences are that we use different flags, a 64 bit LCG
|
|
// and handle negatives differently.
|
|
// The core reason we can't use golang's TempFile is that we must write
|
|
// to the file synchronously, as we need this to persist to disk.
|
|
// We also open it in write-only mode, to avoid concerns that arise with read.
|
|
var (
|
|
dir = filepath.Dir(filename)
|
|
f *os.File
|
|
)
|
|
|
|
nconflict := 0
|
|
// Limit the number of attempts to create a file. Something is seriously
|
|
// wrong if it didn't get created after 1000 attempts, and we don't want
|
|
// an infinite loop
|
|
i := 0
|
|
for ; i < atomicWriteFileMaxNumWriteAttempts; i++ {
|
|
name := filepath.Join(dir, atomicWriteFilePrefix+randWriteFileSuffix())
|
|
f, err = os.OpenFile(name, atomicWriteFileFlag, perm)
|
|
// If the file already exists, try a new file
|
|
if os.IsExist(err) {
|
|
// If the files exists too many times, start reseeding as we've
|
|
// likely hit another instances seed.
|
|
if nconflict++; nconflict > atomicWriteFileMaxNumConflicts {
|
|
atomicWriteFileRandMu.Lock()
|
|
atomicWriteFileRand = writeFileRandReseed()
|
|
atomicWriteFileRandMu.Unlock()
|
|
}
|
|
continue
|
|
} else if err != nil {
|
|
return err
|
|
}
|
|
break
|
|
}
|
|
if i == atomicWriteFileMaxNumWriteAttempts {
|
|
return fmt.Errorf("could not create atomic write file after %d attempts", i)
|
|
}
|
|
|
|
// Clean up in any case. Defer stacking order is last-in-first-out.
|
|
defer os.Remove(f.Name())
|
|
defer f.Close()
|
|
|
|
if n, err := f.Write(data); err != nil {
|
|
return err
|
|
} else if n < len(data) {
|
|
return io.ErrShortWrite
|
|
}
|
|
// Close the file before renaming it, otherwise it will cause "The process
|
|
// cannot access the file because it is being used by another process." on windows.
|
|
f.Close()
|
|
|
|
return os.Rename(f.Name(), filename)
|
|
}
|