Give frozen files their own filesystem.FileReader type

We currently have Close() and ReadAt() operations on the file itself.
This is fairly confusing, considering that those should only be invoked
when the file is frozen.

Author: EdSchouten

pkg/filesystem/virtual/pool_backed_file_allocator.go

@@ -128,29 +128,15 @@ func (f *fileBackedFile) lockMutatingData() {
 	}
 }
 
-func (f *fileBackedFile) acquireFrozenDescriptorLocked() (success bool) {
+func (f *fileBackedFile) acquireFrozenDescriptorLocked() (filesystem.FileReader, bool) {
 	if f.referenceCount == 0 {
-		return false
+		return nil, false
 	}
 	f.referenceCount++
 	f.frozenDescriptorsCount++
-	return true
-}
-
-func (f *fileBackedFile) releaseFrozenDescriptor() {
-	f.lock.Lock()
-	defer f.lock.Unlock()
-
-	if f.frozenDescriptorsCount == 0 {
-		panic("Invalid open frozen descriptor count")
-	}
-	f.frozenDescriptorsCount--
-	if f.frozenDescriptorsCount == 0 && f.unfreezeWakeup != nil {
-		close(f.unfreezeWakeup)
-		f.unfreezeWakeup = nil
-	}
-
-	f.releaseReferencesLocked(1)
+	return &frozenFileBackedFile{
+		file: f,
+	}, true
 }
 
 func (f *fileBackedFile) acquireShareAccessLocked(shareAccess ShareMask) {
@@ -199,15 +185,15 @@ func (f *fileBackedFile) getCachedDigest() digest.Digest {
 // a cached value, or computes the digest and caches it. It is only safe
 // to call this function while the file is frozen (i.e., calling
 // f.acquireFrozenDescriptorLocked()).
-func (f *fileBackedFile) updateCachedDigest(digestFunction digest.Function) (digest.Digest, error) {
+func (f *fileBackedFile) updateCachedDigest(digestFunction digest.Function, frozenFile filesystem.FileReader) (digest.Digest, error) {
 	// Check whether the cached digest we have is still valid.
 	if cachedDigest := f.getCachedDigest(); cachedDigest != digest.BadDigest && cachedDigest.UsesDigestFunction(digestFunction) {
 		return cachedDigest, nil
 	}
 
 	// If not, compute a new digest.
 	digestGenerator := digestFunction.NewGenerator(math.MaxInt64)
-	if _, err := io.Copy(digestGenerator, io.NewSectionReader(f, 0, math.MaxInt64)); err != nil {
+	if _, err := io.Copy(digestGenerator, io.NewSectionReader(frozenFile, 0, math.MaxInt64)); err != nil {
 		return digest.BadDigest, util.StatusWrapWithCode(err, codes.Internal, "Failed to compute file digest")
 	}
 	newDigest := digestGenerator.Sum()
@@ -250,22 +236,22 @@ func (f *fileBackedFile) uploadFile(ctx context.Context, contentAddressableStora
 			poolBackedFileAllocatorWritableFileUploadDelaySeconds.Observe(time.Now().Sub(start).Seconds())
 		}
 	}
-	success := f.acquireFrozenDescriptorLocked()
+	frozenFile, success := f.acquireFrozenDescriptorLocked()
 	f.lock.Unlock()
 	if !success {
 		return digest.BadDigest, status.Error(codes.NotFound, "File was unlinked before uploading could start")
 	}
 
-	blobDigest, err := f.updateCachedDigest(digestFunction)
+	blobDigest, err := f.updateCachedDigest(digestFunction, frozenFile)
 	if err != nil {
-		f.Close()
+		frozenFile.Close()
 		return digest.BadDigest, err
 	}
 
 	if err := contentAddressableStorage.Put(
 		ctx,
 		blobDigest,
-		buffer.NewValidatedBufferFromReaderAt(f, blobDigest.GetSizeBytes())); err != nil {
+		buffer.NewValidatedBufferFromReaderAt(frozenFile, blobDigest.GetSizeBytes())); err != nil {
 		return digest.BadDigest, util.StatusWrap(err, "Failed to upload file")
 	}
 	return blobDigest, nil
@@ -275,14 +261,14 @@ func (f *fileBackedFile) getBazelOutputServiceStat(digestFunction *digest.Functi
 	var locator *anypb.Any
 	f.lock.Lock()
 	if f.writableDescriptorsCount == 0 {
-		success := f.acquireFrozenDescriptorLocked()
+		frozenFile, success := f.acquireFrozenDescriptorLocked()
 		f.lock.Unlock()
 		if !success {
 			return nil, status.Error(codes.NotFound, "File was unlinked before digest computation could start")
 		}
-		defer f.releaseFrozenDescriptor()
+		defer frozenFile.Close()
 
-		blobDigest, err := f.updateCachedDigest(*digestFunction)
+		blobDigest, err := f.updateCachedDigest(*digestFunction, frozenFile)
 		if err != nil {
 			return nil, err
 		}
@@ -312,18 +298,6 @@ func (f *fileBackedFile) getBazelOutputServiceStat(digestFunction *digest.Functi
 	}, nil
 }
 
-func (f *fileBackedFile) Close() error {
-	f.releaseFrozenDescriptor()
-	return nil
-}
-
-func (f *fileBackedFile) ReadAt(b []byte, off int64) (int, error) {
-	f.lock.Lock()
-	defer f.lock.Unlock()
-
-	return f.file.ReadAt(b, off)
-}
-
 func (f *fileBackedFile) VirtualAllocate(off, size uint64) Status {
 	f.lockMutatingData()
 	defer f.lock.Unlock()
@@ -540,3 +514,45 @@ func (f *fileBackedFile) VirtualWrite(buf []byte, offset uint64) (int, Status) {
 	}
 	return nWritten, StatusOK
 }
+
+// frozenFileBackedFile is a handle to a file whose contents have been
+// frozen and has been opened for reading. The file will be unfrozen
+// when closed.
+type frozenFileBackedFile struct {
+	file *fileBackedFile
+}
+
+func (ff *frozenFileBackedFile) Close() error {
+	f := ff.file
+	f.lock.Lock()
+	defer f.lock.Unlock()
+
+	if f.frozenDescriptorsCount == 0 {
+		panic("Invalid open frozen descriptor count")
+	}
+	f.frozenDescriptorsCount--
+	if f.frozenDescriptorsCount == 0 && f.unfreezeWakeup != nil {
+		close(f.unfreezeWakeup)
+		f.unfreezeWakeup = nil
+	}
+
+	f.releaseReferencesLocked(1)
+	ff.file = nil
+	return nil
+}
+
+func (ff *frozenFileBackedFile) GetNextRegionOffset(offset int64, regionType filesystem.RegionType) (int64, error) {
+	f := ff.file
+	f.lock.Lock()
+	defer f.lock.Unlock()
+
+	return f.file.GetNextRegionOffset(offset, regionType)
+}
+
+func (ff *frozenFileBackedFile) ReadAt(b []byte, off int64) (int, error) {
+	f := ff.file
+	f.lock.Lock()
+	defer f.lock.Unlock()
+
+	return f.file.ReadAt(b, off)
+}