@@ -140,20 +140,33 @@ type cacheEntry struct {
140
140
141
141
type fetcher func (context.Context , uuid.UUID ) (CacheEntryValue , error )
142
142
143
+ var _ fs.FS = (* CloseFS )(nil )
144
+
145
+ // CloseFS is a wrapper around fs.FS that implements io.Closer. The Close()
146
+ // method tells the cache to release the fileID. Once all open references are
147
+ // closed, the file is removed from the cache.
148
+ type CloseFS struct {
149
+ fs.FS
150
+
151
+ close func ()
152
+ }
153
+
154
+ func (f * CloseFS ) Close () { f .close () }
155
+
143
156
// Acquire will load the fs.FS for the given file. It guarantees that parallel
144
157
// calls for the same fileID will only result in one fetch, and that parallel
145
158
// calls for distinct fileIDs will fetch in parallel.
146
159
//
147
160
// Safety: Every call to Acquire that does not return an error must have a
148
161
// matching call to Release.
149
- func (c * Cache ) Acquire (ctx context.Context , fileID uuid.UUID ) (fs. FS , error ) {
150
- // It's important that this `Load` call occurs outside of `prepare`, after the
162
+ func (c * Cache ) Acquire (ctx context.Context , fileID uuid.UUID ) (* CloseFS , error ) {
163
+ // It's important that this `Load` call occurs outside `prepare`, after the
151
164
// mutex has been released, or we would continue to hold the lock until the
152
165
// entire file has been fetched, which may be slow, and would prevent other
153
166
// files from being fetched in parallel.
154
167
it , err := c .prepare (ctx , fileID ).Load ()
155
168
if err != nil {
156
- c .Release (fileID )
169
+ c .release (fileID )
157
170
return nil , err
158
171
}
159
172
@@ -163,11 +176,19 @@ func (c *Cache) Acquire(ctx context.Context, fileID uuid.UUID) (fs.FS, error) {
163
176
}
164
177
// Always check the caller can actually read the file.
165
178
if err := c .authz .Authorize (ctx , subject , policy .ActionRead , it .Object ); err != nil {
166
- c .Release (fileID )
179
+ c .release (fileID )
167
180
return nil , err
168
181
}
169
182
170
- return it .FS , err
183
+ var once sync.Once
184
+ return & CloseFS {
185
+ FS : it .FS ,
186
+ close : func () {
187
+ // sync.Once makes the Close() idempotent, so we can call it
188
+ // multiple times without worrying about double-releasing.
189
+ once .Do (func () { c .release (fileID ) })
190
+ },
191
+ }, nil
171
192
}
172
193
173
194
func (c * Cache ) prepare (ctx context.Context , fileID uuid.UUID ) * lazy.ValueWithError [CacheEntryValue ] {
@@ -203,9 +224,12 @@ func (c *Cache) prepare(ctx context.Context, fileID uuid.UUID) *lazy.ValueWithEr
203
224
return entry .value
204
225
}
205
226
206
- // Release decrements the reference count for the given fileID, and frees the
227
+ // release decrements the reference count for the given fileID, and frees the
207
228
// backing data if there are no further references being held.
208
- func (c * Cache ) Release (fileID uuid.UUID ) {
229
+ //
230
+ // release should only be called after a successful call to Acquire using the Release()
231
+ // method on the returned *CloseFS.
232
+ func (c * Cache ) release (fileID uuid.UUID ) {
209
233
c .lock .Lock ()
210
234
defer c .lock .Unlock ()
211
235
0 commit comments