@@ -2,8 +2,10 @@ package ioutils
2
2
3
3
import (
4
4
"io"
5
+ "math/rand"
5
6
"os"
6
7
"path/filepath"
8
+ "strconv"
7
9
"time"
8
10
)
9
11
@@ -17,6 +19,9 @@ type AtomicFileWriterOptions struct {
17
19
// On successful return from Close() this is set to the mtime of the
18
20
// newly written file.
19
21
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
20
25
}
21
26
22
27
var defaultWriterOptions = AtomicFileWriterOptions {}
@@ -38,22 +43,33 @@ func NewAtomicFileWriterWithOpts(filename string, perm os.FileMode, opts *Atomic
38
43
// temporary file and closing it atomically changes the temporary file to
39
44
// destination path. Writing and closing concurrently is not allowed.
40
45
func 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 )
45
48
if opts == nil {
46
49
opts = & defaultWriterOptions
47
50
}
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
+ }
48
63
abspath , err := filepath .Abs (filename )
49
64
if err != nil {
50
65
return nil , err
51
66
}
52
67
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 ,
57
73
}, nil
58
74
}
59
75
@@ -91,12 +107,13 @@ func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error {
91
107
}
92
108
93
109
type 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
100
117
}
101
118
102
119
func (w * atomicFileWriter ) Write (dt []byte ) (int , error ) {
@@ -108,22 +125,35 @@ func (w *atomicFileWriter) Write(dt []byte) (int, error) {
108
125
}
109
126
110
127
func (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
114
142
}
115
- }()
143
+ }
116
144
if ! w .noSync {
117
145
if err := fdatasync (w .f ); err != nil {
118
146
w .f .Close ()
119
147
return err
120
148
}
121
149
}
122
150
151
+ var size int64 = 0
123
152
// fstat before closing the fd
124
153
info , statErr := w .f .Stat ()
125
154
if statErr == nil {
126
155
w .modTime = info .ModTime ()
156
+ size = info .Size ()
127
157
}
128
158
// We delay error reporting until after the real call to close()
129
159
// to match the traditional linux close() behaviour that an fd
@@ -138,11 +168,51 @@ func (w *atomicFileWriter) Close() (retErr error) {
138
168
return statErr
139
169
}
140
170
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 {
142
173
return err
143
174
}
144
175
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
+ } else {
185
+ return os .Rename (tmpName , w .fn )
186
+ }
187
+ }
188
+ return nil
189
+ }
190
+
191
+ // ensure that the file is of at least indicated size
192
+ func preAllocate (filename string , size int64 , perm os.FileMode ) error {
193
+ f , err := os .OpenFile (filename , os .O_WRONLY | os .O_CREATE | os .O_APPEND , perm )
194
+ if err != nil {
195
+ return err
196
+ }
197
+ defer f .Close ()
198
+ info , err := f .Stat ()
199
+ if err != nil {
200
+ return err
201
+ }
202
+ extendBytes := size - info .Size ()
203
+ if extendBytes > 0 {
204
+ var blockSize int64 = 65536
205
+ block := make ([]byte , blockSize )
206
+ for extendBytes > 0 {
207
+ if blockSize > extendBytes {
208
+ blockSize = extendBytes
209
+ }
210
+ _ , err := f .Write (block [:blockSize ])
211
+ if err != nil {
212
+ return err
213
+ }
214
+ extendBytes -= blockSize
215
+ }
146
216
}
147
217
return nil
148
218
}
0 commit comments