Fix race conditions in file cache and block cache

component/block_cache/block_cache.go

@@ -482,6 +482,17 @@ func (bc *BlockCache) OpenFile(options internal.OpenFileOptions) (*handlemap.Han
 		if !valid {
 			return nil, fmt.Errorf("block size mismatch for %s", options.Name)
 		}
+
+		// Size and block list are fetched separately, so a concurrent write or truncate can make them disagree.
+		if n := len(*blockList); n > 0 && int64(n-1)*int64(bc.blockSize) >= handle.Size {
+			log.Err(
+				"BlockCache::OpenFile : Block list of %s does not match size %v [blocks: %v]",
+				options.Name,
+				handle.Size,
+				n,
+			)
+			return nil, fmt.Errorf("block list does not match size for %s", options.Name)
+		}
 	} else if options.Flags&os.O_WRONLY != 0 {
 		handle.RemoveValue("ETAG")
 	}
@@ -667,17 +678,17 @@ func (bc *BlockCache) getBlockSize(fileSize uint64, block *Block) uint64 {
 
 // ReadInBuffer: Read the file into a buffer
 func (bc *BlockCache) ReadInBuffer(options *internal.ReadInBufferOptions) (int, error) {
-	if options.Offset >= options.Handle.Size {
-		// EOF reached so early exit
-		return 0, io.EOF
-	}
-
 	// As of now we allow only one operation on a handle at a time
 	// This simplifies the logic of block-cache otherwise we will have to handle
 	// a lot of race conditions and logic becomes complex and sub-performant
 	options.Handle.Lock()
 	defer options.Handle.Unlock()
 
+	if options.Offset >= options.Handle.Size {
+		// EOF reached so early exit
+		return 0, io.EOF
+	}
+
 	// Keep getting next blocks until you read the request amount of data
 	dataRead := int(0)
 	for dataRead < len(options.Data) {
@@ -846,15 +857,15 @@ func (bc *BlockCache) getBlock(handle *handlemap.Handle, readoffset uint64) (*Bl
 			}
 		} else {
 			// This is a case of random read so increment the random read count
-			handle.OptCnt++
+			optCnt := handle.OptCnt.Add(1)
 
 			log.Debug(
 				"BlockCache::getBlock : Unable to get block %v=>%s (offset %v, index %v) Random %v",
 				handle.ID,
 				handle.Path,
 				readoffset,
 				index,
-				handle.OptCnt,
+				optCnt,
 			)
 
 			// This block is not present even after prefetch so lets download it now
@@ -908,7 +919,7 @@ func (bc *BlockCache) getBlock(handle *handlemap.Handle, readoffset uint64) (*Bl
 			block.flags.Clear(BlockFlagDownloading)
 
 			// Download complete and you are first reader of this block
-			if !bc.noPrefetch && handle.OptCnt <= MIN_RANDREAD {
+			if !bc.noPrefetch && handle.OptCnt.Load() <= MIN_RANDREAD {
 				// So far this file has been read sequentially so prefetch more
 				val, _ := handle.GetValue("#")
 				if int64(val.(uint64)*bc.blockSize) < handle.Size {
@@ -976,7 +987,7 @@ func (bc *BlockCache) startPrefetch(handle *handlemap.Handle, index uint64, pref
 	currentCnt := handle.Buffers.Cooked.Len() + handle.Buffers.Cooking.Len()
 	cnt := uint32(0)
 
-	if handle.OptCnt > MIN_RANDREAD {
+	if handle.OptCnt.Load() > MIN_RANDREAD {
 		// This handle has been read randomly and we have reached the threshold to declare a random read case
 
 		if currentCnt > MIN_PREFETCH {
@@ -1190,6 +1201,7 @@ func (bc *BlockCache) lineupDownload(handle *handlemap.Handle, block *Block, pre
 		failCnt:  0,
 		upload:   false,
 		ETag:     Etag,
+		fileSize: handle.Size,
 	}
 
 	// Remove this block from free block list and add to in-process list
@@ -1258,12 +1270,12 @@ func (bc *BlockCache) download(item *workItem) {
 				}
 
 				if numberOfBytes != int(bc.blockSize) &&
-					item.block.offset+uint64(numberOfBytes) != uint64(item.handle.Size) {
+					item.block.offset+uint64(numberOfBytes) != uint64(item.fileSize) {
 					log.Err(
 						"BlockCache::download : Local data retrieved from disk size mismatch, Expected %v, OnDisk %v, fileSize %v",
-						bc.getBlockSize(uint64(item.handle.Size), item.block),
+						bc.getBlockSize(uint64(item.fileSize), item.block),
 						numberOfBytes,
-						item.handle.Size,
+						item.fileSize,
 					)
 					successfulRead = false
 					f.Close()
@@ -1305,8 +1317,7 @@ func (bc *BlockCache) download(item *workItem) {
 	if item.failCnt > MAX_FAIL_CNT {
 		// If we failed to read the data 3 times then just give up
 		log.Err(
-			"BlockCache::download : 3 attempts to download a block have failed %v=>%s (index %v, offset %v)",
-			item.handle.ID,
+			"BlockCache::download : 3 attempts to download a block have failed %s (index %v, offset %v)",
 			item.handle.Path,
 			item.block.id,
 			item.block.offset,
@@ -1319,8 +1330,7 @@ func (bc *BlockCache) download(item *workItem) {
 	if err != nil && err != io.EOF {
 		// Fail to read the data so just reschedule this request
 		log.Err(
-			"BlockCache::download : Failed to read %v=>%s from offset %v [%s]",
-			item.handle.ID,
+			"BlockCache::download : Failed to read %s from offset %v [%s]",
 			item.handle.Path,
 			item.block.id,
 			err.Error(),
@@ -1331,8 +1341,7 @@ func (bc *BlockCache) download(item *workItem) {
 	} else if n == 0 {
 		// No data read so just reschedule this request
 		log.Err(
-			"BlockCache::download : Failed to read %v=>%s from offset %v [0 bytes read]",
-			item.handle.ID,
+			"BlockCache::download : Failed to read %s from offset %v [0 bytes read]",
 			item.handle.Path,
 			item.block.id,
 		)
@@ -1345,8 +1354,7 @@ func (bc *BlockCache) download(item *workItem) {
 	if etag != "" {
 		if item.ETag != "" && item.ETag != etag {
 			log.Err(
-				"BlockCache::download : Blob has changed for %v=>%s (index %v, offset %v)",
-				item.handle.ID,
+				"BlockCache::download : Blob has changed for %s (index %v, offset %v)",
 				item.handle.Path,
 				item.block.id,
 				item.block.offset,
@@ -1600,7 +1608,8 @@ func (bc *BlockCache) getOrCreateBlock(handle *handlemap.Handle, offset uint64)
 		block = node.(*Block)
 
 		// If the block was staged earlier then we are overwriting it here so move it back to cooking queue
-		if block.flags.IsSet(BlockFlagSynced) {
+		// The upload worker sets Synced before it signals completion, so wait for it if it is still uploading
+		if block.flags.IsSet(BlockFlagSynced) && !block.flags.IsSet(BlockFlagUploading) {
 			log.Debug(
 				"BlockCache::getOrCreateBlock : Overwriting back to staged block %v for %v=>%s",
 				block.id,
@@ -1777,6 +1786,7 @@ func (bc *BlockCache) lineupUpload(
 		failCnt:  0,
 		upload:   true,
 		blockId:  id,
+		fileSize: handle.Size,
 	}
 
 	block.Uploading()
@@ -1854,7 +1864,7 @@ func (bc *BlockCache) upload(item *workItem) {
 	flock := bc.fileLocks.Get(fileName)
 	flock.Lock()
 	defer flock.Unlock()
-	blockSize := bc.getBlockSize(uint64(item.handle.Size), item.block)
+	blockSize := bc.getBlockSize(uint64(item.fileSize), item.block)
 	// This block is updated so we need to stage it now
 	err := bc.NextComponent().StageData(internal.StageDataOptions{
 		Name:   item.handle.Path,

component/block_cache/block_cache_test.go

@@ -1445,7 +1445,8 @@ func (suite *blockCacheTestSuite) TestZZZZLazyWrite() {
 	suite.assert.True(handle.Dirty())
 
 	_ = tobj.blockCache.ReleaseFile(internal.ReleaseFileOptions{Handle: handle})
-	time.Sleep(1 * time.Second)
+	// wait for the async close before turning lazy write back off
+	tobj.blockCache.fileCloseOpt.Wait()
 	tobj.blockCache.lazyWrite = false
 
 	// As lazy write is enabled flush shall not upload the file
@@ -3320,6 +3321,121 @@ func (suite *blockCacheTestSuite) TestReadCommittedLastBlocksOverwrite() {
 	suite.assert.Equal(h.Size, int64((15*_1MB)+(_1MB/2)))
 }
 
+// An overwrite must wait for an in-flight upload of the same block. The upload worker marks
+// the block Synced before it clears the dirty bit and signals completion, so a writer that
+// trusts Synced alone can have its dirty bit wiped and its data never uploaded.
+func (suite *blockCacheTestSuite) TestOverwriteWaitsForInFlightUpload() {
+	cfg := "block_cache:\n  block-size-mb: 1\n  mem-size-mb: 20\n  prefetch: 12\n  parallelism: 10"
+	tobj, err := setupPipeline(cfg)
+	defer tobj.cleanupPipeline()
+	suite.assert.NoError(err)
+
+	path := getTestFileName(suite.T().Name())
+	h, err := tobj.blockCache.CreateFile(internal.CreateFileOptions{Name: path, Mode: 0777})
+	suite.assert.NoError(err)
+
+	_, err = tobj.blockCache.WriteFile(
+		&internal.WriteFileOptions{Handle: h, Offset: 0, Data: dataBuff[:10]},
+	)
+	suite.assert.NoError(err)
+	node, found := h.GetValue("0")
+	suite.assert.True(found)
+	block := node.(*Block)
+
+	// Put the block in the state an upload worker leaves it in just before it finishes:
+	// queued for upload and already marked Synced, but still dirty and not yet signalled.
+	h.Lock()
+	block.Uploading()
+	block.flags.Set(BlockFlagUploading)
+	block.flags.Set(BlockFlagSynced)
+	tobj.blockCache.addToCooked(h, block)
+	h.Unlock()
+
+	done := make(chan error, 1)
+	go func() {
+		_, werr := tobj.blockCache.WriteFile(
+			&internal.WriteFileOptions{Handle: h, Offset: 0, Data: dataBuff[10:30]},
+		)
+		done <- werr
+	}()
+
+	select {
+	case <-done:
+		suite.assert.Fail("overwrite did not wait for the in-flight upload to finish")
+		return
+	case <-time.After(200 * time.Millisecond):
+	}
+
+	// Let the simulated upload worker finish
+	block.NoMoreDirty()
+	block.Ready(BlockStatusUploaded)
+
+	select {
+	case err = <-done:
+		suite.assert.NoError(err)
+	case <-time.After(5 * time.Second):
+		suite.assert.Fail("overwrite never completed")
+	}
+	suite.assert.True(block.IsDirty(), "overwrite lost its dirty bit")
+
+	err = tobj.blockCache.ReleaseFile(internal.ReleaseFileOptions{Handle: h})
+	suite.assert.NoError(err)
+
+	data, err := os.ReadFile(filepath.Join(tobj.fake_storage_path, path))
+	suite.assert.NoError(err)
+	suite.assert.Equal(dataBuff[10:30], data)
+}
+
+// An upload must stage the block at the size recorded in the block list when it was lined up.
+// If it reads the live handle size instead, a concurrent write that extends the file makes it
+// stage a full block while the block list still records the short size, and the commit then
+// pads it with a filler block, corrupting the file.
+func (suite *blockCacheTestSuite) TestUploadRacingWriteThatExtendsFile() {
+	cfg := "block_cache:\n  block-size-mb: 1\n  mem-size-mb: 20\n  prefetch: 12\n  parallelism: 10"
+	tobj, err := setupPipeline(cfg)
+	defer tobj.cleanupPipeline()
+	suite.assert.NoError(err)
+
+	for i := range 20 {
+		path := fmt.Sprintf("%s_%d", getTestFileName(suite.T().Name()), i)
+		h, err := tobj.blockCache.CreateFile(internal.CreateFileOptions{Name: path, Mode: 0777})
+		suite.assert.NoError(err)
+
+		_, err = tobj.blockCache.WriteFile(
+			&internal.WriteFileOptions{Handle: h, Offset: 0, Data: dataBuff[:10]},
+		)
+		suite.assert.NoError(err)
+
+		h.Lock()
+		err = tobj.blockCache.stageBlocks(h, 1)
+		lst, _ := h.GetValue("blockList")
+		staged := *lst.(map[int64]*blockInfo)[0]
+		h.Unlock()
+		suite.assert.NoError(err)
+
+		// extend the file while block 0 may still be uploading
+		_, err = tobj.blockCache.WriteFile(
+			&internal.WriteFileOptions{Handle: h, Offset: int64(2 * _1MB), Data: dataBuff[:10]},
+		)
+		suite.assert.NoError(err)
+
+		// wait for the upload of block 0, then compare what was staged with what was recorded
+		h.Lock()
+		tobj.blockCache.waitAndFreeUploadedBlocks(h, 1)
+		h.Unlock()
+		stagedPath := filepath.Join(tobj.fake_storage_path, path) + "_" +
+			strings.ReplaceAll(staged.id, "/", "_")
+		fi, err := os.Stat(stagedPath)
+		suite.assert.NoError(err)
+		if err == nil {
+			suite.assert.Equal(int64(staged.size), fi.Size(), "iteration %d", i)
+		}
+
+		err = tobj.blockCache.ReleaseFile(internal.ReleaseFileOptions{Handle: h})
+		suite.assert.NoError(err)
+	}
+}
+
 // In order for 'go test' to run this suite, we need to create
 // a normal test function and pass our suite to suite.Run
 func TestBlockCacheTestSuite(t *testing.T) {

component/block_cache/threadpool.go

@@ -62,6 +62,7 @@ type workItem struct {
 	upload   bool              // Flag marking this is a upload request or not
 	blockId  string            // BlockId of the block
 	ETag     string            // Etag of the file before scheduling.
+	fileSize int64             // Size of the file when this item was scheduled
 }
 
 // Reason for storing Etag in workitem struct:

component/block_cache/threadpool_test.go

@@ -126,8 +126,9 @@ func (suite *threadPoolTestSuite) TestPrioritySchedule() {
 		tp.Schedule(i < 20, &workItem{failCnt: 5})
 	}
 
-	time.Sleep(100 * time.Millisecond)
-	suite.assert.Equal(int32(100), callbackCnt)
+	suite.assert.Eventually(func() bool {
+		return atomic.LoadInt32(&callbackCnt) == 100
+	}, 5*time.Second, 10*time.Millisecond)
 	tp.Stop()
 }
 
@@ -158,9 +159,10 @@ func (suite *threadPoolTestSuite) TestPriorityScheduleWithWriter() {
 		tp.Schedule(i < 20, &workItem{failCnt: 5, upload: true, blockId: "test"})
 	}
 
-	time.Sleep(100 * time.Millisecond)
-	suite.assert.Equal(int32(100), callbackWCnt)
-	suite.assert.Equal(int32(0), callbackRCnt)
+	suite.assert.Eventually(func() bool {
+		return atomic.LoadInt32(&callbackWCnt) == 100
+	}, 5*time.Second, 10*time.Millisecond)
+	suite.assert.Equal(int32(0), atomic.LoadInt32(&callbackRCnt))
 	tp.Stop()
 }
 

component/file_cache/async.go

@@ -122,6 +122,7 @@ func (fc *FileCache) configureScheduler() error {
 		// add cron callback
 		entryId, err := fc.cronScheduler.AddFunc(window.cronExpr, func() {
 			// Is this a transition from inactive?
+			fc.scheduleLock.Lock()
 			windowCount := fc.activeWindows.Add(1)
 			if windowCount == 1 {
 				// transition to active - open the window
@@ -131,6 +132,7 @@ func (fc *FileCache) configureScheduler() error {
 					window.name,
 				)
 			}
+			fc.scheduleLock.Unlock()
 			log.Info(
 				"FileCache::SchedulerCronFunc : %s (%s) started (numActive=%d)",
 				window.name,
@@ -153,6 +155,7 @@ func (fc *FileCache) configureScheduler() error {
 					return
 				case <-ctx.Done():
 					// Window has completed, update active window count
+					fc.scheduleLock.Lock()
 					windowCount = fc.activeWindows.Add(-1)
 					log.Info(
 						"FileCache::SchedulerCronFunc : %s (%s) ended (numActive=%d)",
@@ -168,6 +171,7 @@ func (fc *FileCache) configureScheduler() error {
 							window.name,
 						)
 					}
+					fc.scheduleLock.Unlock()
 					return
 				}
 			}
@@ -212,6 +216,13 @@ func (fc *FileCache) startScheduler() {
 	fc.cronScheduler.Start()
 }
 
+// uploadWindow returns a channel that is closed while scheduled uploads are allowed
+func (fc *FileCache) uploadWindow() <-chan struct{} {
+	fc.scheduleLock.RLock()
+	defer fc.scheduleLock.RUnlock()
+	return fc.startScheduledUploads
+}
+
 // flock must be locked
 func (fc *FileCache) addPendingOp(name string, value pendingFlags) {
 	log.Trace("FileCache::addPendingOp : %s", name)
@@ -237,7 +248,7 @@ func (fc *FileCache) servicePendingOps() {
 		case <-fc.componentStopping:
 			log.Crit("FileCache::servicePendingOps : Stopping")
 			return
-		case <-fc.startScheduledUploads:
+		case <-fc.uploadWindow():
 			if retryDelay > 0 {
 				select {
 				case <-time.After(retryDelay):
@@ -293,7 +304,7 @@ func (fc *FileCache) runPendingOpCycle() (int, error) {
 		select {
 		case <-fc.componentStopping:
 			return false
-		case <-fc.startScheduledUploads:
+		case <-fc.uploadWindow():
 			name := key.(string)
 			numFilesProcessed++
 			if !fc.updateObject(name, value.(pendingFlags)) {