@@ -2,8 +2,10 @@ package ioutils
22
33import (
44 "io"
5+ "math/rand"
56 "os"
67 "path/filepath"
8+ "strconv"
79 "time"
810)
911
@@ -17,6 +19,9 @@ type AtomicFileWriterOptions struct {
1719 // On successful return from Close() this is set to the mtime of the
1820 // newly written file.
1921 ModTime time.Time
22+ // Whenever an atomic file is successfully written create a temporary
23+ // file of the same size in order to pre-allocate storage for the next write
24+ PreAllocate bool
2025}
2126
2227var defaultWriterOptions = AtomicFileWriterOptions {}
@@ -38,22 +43,33 @@ func NewAtomicFileWriterWithOpts(filename string, perm os.FileMode, opts *Atomic
3843// temporary file and closing it atomically changes the temporary file to
3944// destination path. Writing and closing concurrently is not allowed.
4045func newAtomicFileWriter (filename string , perm os.FileMode , opts * AtomicFileWriterOptions ) (* atomicFileWriter , error ) {
41- f , err := os .CreateTemp (filepath .Dir (filename ), ".tmp-" + filepath .Base (filename ))
42- if err != nil {
43- return nil , err
44- }
46+ dir := filepath .Dir (filename )
47+ base := filepath .Base (filename )
4548 if opts == nil {
4649 opts = & defaultWriterOptions
4750 }
51+ random := ""
52+ // need predictable name when pre-allocated
53+ if ! opts .PreAllocate {
54+ random = strconv .FormatUint (rand .Uint64 (), 36 )
55+ }
56+ tmp := filepath .Join (dir , ".tmp" + random + "-" + base )
57+ // if pre-allocated the temporary file exists and contains some data
58+ // do not truncate it here, instead truncate at Close()
59+ f , err := os .OpenFile (tmp , os .O_WRONLY | os .O_CREATE , perm )
60+ if err != nil {
61+ return nil , err
62+ }
4863 abspath , err := filepath .Abs (filename )
4964 if err != nil {
5065 return nil , err
5166 }
5267 return & atomicFileWriter {
53- f : f ,
54- fn : abspath ,
55- perm : perm ,
56- noSync : opts .NoSync ,
68+ f : f ,
69+ fn : abspath ,
70+ perm : perm ,
71+ noSync : opts .NoSync ,
72+ preAllocate : opts .PreAllocate ,
5773 }, nil
5874}
5975
@@ -91,12 +107,13 @@ func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error {
91107}
92108
93109type atomicFileWriter struct {
94- f * os.File
95- fn string
96- writeErr error
97- perm os.FileMode
98- noSync bool
99- modTime time.Time
110+ f * os.File
111+ fn string
112+ writeErr error
113+ perm os.FileMode
114+ noSync bool
115+ preAllocate bool
116+ modTime time.Time
100117}
101118
102119func (w * atomicFileWriter ) Write (dt []byte ) (int , error ) {
@@ -108,22 +125,35 @@ func (w *atomicFileWriter) Write(dt []byte) (int, error) {
108125}
109126
110127func (w * atomicFileWriter ) Close () (retErr error ) {
111- defer func () {
112- if retErr != nil || w .writeErr != nil {
113- os .Remove (w .f .Name ())
128+ if ! w .preAllocate {
129+ defer func () {
130+ if retErr != nil || w .writeErr != nil {
131+ os .Remove (w .f .Name ())
132+ }
133+ }()
134+ } else {
135+ truncateAt , err := w .f .Seek (0 , io .SeekCurrent )
136+ if err != nil {
137+ return err
138+ }
139+ err = w .f .Truncate (truncateAt )
140+ if err != nil {
141+ return err
114142 }
115- }()
143+ }
116144 if ! w .noSync {
117145 if err := fdatasync (w .f ); err != nil {
118146 w .f .Close ()
119147 return err
120148 }
121149 }
122150
151+ var size int64 = 0
123152 // fstat before closing the fd
124153 info , statErr := w .f .Stat ()
125154 if statErr == nil {
126155 w .modTime = info .ModTime ()
156+ size = info .Size ()
127157 }
128158 // We delay error reporting until after the real call to close()
129159 // to match the traditional linux close() behaviour that an fd
@@ -138,11 +168,50 @@ func (w *atomicFileWriter) Close() (retErr error) {
138168 return statErr
139169 }
140170
141- if err := os .Chmod (w .f .Name (), w .perm ); err != nil {
171+ tmpName := w .f .Name ()
172+ if err := os .Chmod (tmpName , w .perm ); err != nil {
142173 return err
143174 }
144175 if w .writeErr == nil {
145- return os .Rename (w .f .Name (), w .fn )
176+ if w .preAllocate {
177+ err := swapOrMove (tmpName , w .fn )
178+ if err != nil {
179+ return err
180+ }
181+ // ignore errors, this is a best effort operation
182+ preAllocate (tmpName , size , w .perm )
183+ return nil
184+ }
185+ return os .Rename (tmpName , w .fn )
186+ }
187+ return nil
188+ }
189+
190+ // ensure that the file is of at least indicated size
191+ func preAllocate (filename string , size int64 , perm os.FileMode ) error {
192+ f , err := os .OpenFile (filename , os .O_WRONLY | os .O_CREATE | os .O_APPEND , perm )
193+ if err != nil {
194+ return err
195+ }
196+ defer f .Close ()
197+ info , err := f .Stat ()
198+ if err != nil {
199+ return err
200+ }
201+ extendBytes := size - info .Size ()
202+ if extendBytes > 0 {
203+ var blockSize int64 = 65536
204+ block := make ([]byte , blockSize )
205+ for extendBytes > 0 {
206+ if blockSize > extendBytes {
207+ blockSize = extendBytes
208+ }
209+ _ , err := f .Write (block [:blockSize ])
210+ if err != nil {
211+ return err
212+ }
213+ extendBytes -= blockSize
214+ }
146215 }
147216 return nil
148217}
0 commit comments