mrd_instance.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcsx

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"

	"github.com/googlecloudplatform/gcsfuse/v3/cfg"
	"github.com/googlecloudplatform/gcsfuse/v3/internal/cache/lru"
	"github.com/googlecloudplatform/gcsfuse/v3/internal/logger"
	"github.com/googlecloudplatform/gcsfuse/v3/internal/monitor"
	"github.com/googlecloudplatform/gcsfuse/v3/internal/storage/gcs"
	"github.com/googlecloudplatform/gcsfuse/v3/metrics"
	"github.com/jacobsa/fuse/fuseops"
)

const mrdPoolCloseTimeout = 120 * time.Second

// MrdInstance manages a pool of Multi-Range Downloader (MRD) instances for a
// single file inode. It handles the lifecycle of the MRD pool, including
// creation, destruction, and caching.
type MrdInstance struct {
	// mrdPool holds the pool of MultiRangeDownloader instances.
	mrdPool *MRDPool
	// inodeId is the ID of the file inode associated with this instance.
	inodeId fuseops.InodeID
	// object is the GCS object for which the downloaders are created.
	object *gcs.MinObject
	// bucket is the GCS bucket containing the object.
	bucket gcs.Bucket
	// refCount tracks the number of active users of this instance.
	refCount int64
	// refCountMu protects access to refCount.
	refCountMu sync.Mutex
	// poolMu protects access to mrdPool.
	poolMu sync.RWMutex
	// mrdCache is a shared cache for inactive MrdInstance objects.
	mrdCache *lru.Cache
	// holds config specified by the user using config-file flag and CLI flags.
	config *cfg.Config
}

// NewMrdInstance creates a new MrdInstance for a given GCS object.
func NewMrdInstance(obj *gcs.MinObject, bucket gcs.Bucket, cache *lru.Cache, inodeId fuseops.InodeID, config *cfg.Config) *MrdInstance {
	mrdInstance := MrdInstance{
		bucket:   bucket,
		mrdCache: cache,
		inodeId:  inodeId,
		config:   config,
		object:   obj,
	}
	return &mrdInstance
}

// SetMinObject sets the gcs.MinObject stored in the MrdInstance to passed value, only if it's non nil.
// If the generation of the object has changed, it recreates the MRD pool to ensure consistency.
func (mi *MrdInstance) SetMinObject(minObj *gcs.MinObject) error {
	if minObj == nil {
		return fmt.Errorf("MrdInstance::SetMinObject: Missing MinObject")
	}

	mi.poolMu.Lock()
	defer mi.poolMu.Unlock()

	oldObj := mi.object
	// If generation matches, just update the object (e.g. for size updates) and return.
	if oldObj != nil && oldObj.Generation == minObj.Generation {
		mi.object = minObj
		return nil
	}

	// Generations differ, need to create and swap a new pool.
	if err := mi.createAndSwapPool(minObj); err != nil {
		return fmt.Errorf("MrdInstance::SetMinObject: failed to create and swap pool: %w", err)
	}

	return nil
}

// createAndSwapPool creates a new MRD pool and swaps it with the existing one.
// It also updates the minObject with the passed object & closes the old pool.
// LOCKS_REQUIRED(mi.poolMu).
func (mi *MrdInstance) createAndSwapPool(obj *gcs.MinObject) error {
	newPool, err := NewMRDPool(&MRDPoolConfig{PoolSize: int(mi.config.Mrd.PoolSize), object: obj, bucket: mi.bucket}, nil)
	if err != nil {
		return err
	}

	oldPool := mi.mrdPool
	mi.mrdPool = newPool
	mi.object = obj

	closePoolWithTimeout(oldPool, "MrdInstance::createAndSwapPool", mrdPoolCloseTimeout)
	return nil
}

// GetMinObject returns the gcs.MinObject stored in MrdInstance. Used only for unit testing.
func (mi *MrdInstance) GetMinObject() *gcs.MinObject {
	mi.poolMu.RLock()
	defer mi.poolMu.RUnlock()
	return mi.object
}

// getMRDEntry returns a valid MRDEntry from the pool.
// It ensures the pool is initialized and the returned entry is in a usable state,
// recreating it if necessary.
func (mi *MrdInstance) getMRDEntry() (*MRDEntry, error) {
	// Ensure the pool is initialized.
	if err := mi.ensureMRDPool(); err != nil {
		return nil, err
	}

	mi.poolMu.RLock()
	defer mi.poolMu.RUnlock()
	if mi.mrdPool == nil {
		return nil, fmt.Errorf("MrdInstance::getMRDEntry: MRDPool is nil")
	}

	entry := mi.mrdPool.Next()

	// Check if the entry is valid.
	entry.mu.RLock()
	isValid := entry.mrd != nil && entry.mrd.Error() == nil
	entry.mu.RUnlock()

	// Recreate the entry if it is not valid.
	if !isValid {
		if err := mi.mrdPool.RecreateMRD(entry, nil); err != nil {
			return nil, fmt.Errorf("MrdInstance::getMRDEntry: failed to recreate MRD: %w", err)
		}
	}
	return entry, nil
}

// Read downloads data from the GCS object into the provided buffer starting at the offset.
// It handles the details of selecting a valid MRD entry, locking it, and waiting for the async download to complete.
func (mi *MrdInstance) Read(ctx context.Context, p []byte, offset int64, metrics metrics.MetricHandle) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}

	entry, err := mi.getMRDEntry()
	if err != nil {
		return 0, err
	}

	// Prepare the buffer for the read operation.
	// bytes.NewBuffer creates a buffer using p as the initial content.
	// Reset() sets the length to 0 but keeps the capacity, allowing writes to fill p.
	buffer := bytes.NewBuffer(p)
	buffer.Reset()
	done := make(chan readResult, 1)

	// Take Read lock before using MRD for reading.
	entry.mu.RLock()
	if entry.mrd == nil {
		entry.mu.RUnlock()
		return 0, fmt.Errorf("MrdInstance::Read: mrd is nil")
	}

	start := time.Now()
	defer monitor.CaptureMultiRangeDownloaderMetrics(ctx, metrics, "MultiRangeDownloader::Add", start)
	entry.mrd.Add(buffer, offset, int64(len(p)), func(offsetAddCallback int64, bytesReadAddCallback int64, e error) {
		done <- readResult{bytesRead: int(bytesReadAddCallback), err: e}
	})
	entry.mu.RUnlock()

	if !mi.config.FileSystem.IgnoreInterrupts {
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		case res := <-done:
			if res.err != nil && res.err != io.EOF {
				return res.bytesRead, fmt.Errorf("Error in Add call: %w", res.err)
			}
			return res.bytesRead, res.err
		}
	} else {
		res := <-done
		if res.err != nil && res.err != io.EOF {
			return res.bytesRead, fmt.Errorf("Error in Add call: %w", res.err)
		}
		return res.bytesRead, res.err
	}
}

// ensureMRDPool ensures that the MRD pool is initialized. If the pool
// already exists, this function is a no-op.
func (mi *MrdInstance) ensureMRDPool() (err error) {
	// Return early if pool exists.
	mi.poolMu.RLock()
	if mi.mrdPool != nil {
		mi.poolMu.RUnlock()
		return
	}
	mi.poolMu.RUnlock()

	mi.poolMu.Lock()
	defer mi.poolMu.Unlock()

	// Re-check under write lock to handle race condition.
	if mi.mrdPool != nil {
		return
	}

	// Creating a new pool. Not reusing any handle while creating a new pool.
	mi.mrdPool, err = NewMRDPool(&MRDPoolConfig{PoolSize: int(mi.config.Mrd.PoolSize), object: mi.object, bucket: mi.bucket}, nil)
	if err != nil {
		err = fmt.Errorf("MrdInstance::ensureMRDPool Error in creating MRDPool: %w", err)
	}
	return
}

// RecreateMRD recreates the entire MRD pool. This is typically called by the
// file inode when the backing GCS object's generation changes, invalidating
// all existing downloader instances.
func (mi *MrdInstance) RecreateMRD() error {
	mi.poolMu.Lock()
	defer mi.poolMu.Unlock()

	obj := mi.object
	if obj == nil {
		return fmt.Errorf("MrdInstance::RecreateMRD: object is nil")
	}

	if err := mi.createAndSwapPool(obj); err != nil {
		return fmt.Errorf("MrdInstance::RecreateMRD: failed to create new pool: %w", err)
	}

	return nil
}

// closePool closes all MRD instances in the pool and releases associated resources.
func (mi *MrdInstance) closePool() {
	mi.poolMu.Lock()
	defer mi.poolMu.Unlock()
	pool := mi.mrdPool
	mi.mrdPool = nil
	closePoolWithTimeout(pool, "MrdInstance::closePool", mrdPoolCloseTimeout)
}

// closePoolWithTimeout closes the given MRD pool in a separate goroutine with a timeout.
// If closing the pool takes longer than the specified timeout, a warning is logged.
func closePoolWithTimeout(pool *MRDPool, caller string, timeout time.Duration) {
	if pool == nil {
		return
	}

	go func() {
		done := make(chan struct{})
		go func() {
			defer close(done)
			pool.Close()
		}()

		select {
		case <-done:
		case <-time.After(timeout):
			var objectName string
			if pool.poolConfig != nil && pool.poolConfig.object != nil {
				objectName = pool.poolConfig.object.Name
			}
			logger.Warnf("%s: MRDPool.Close() timed out after %v for object %s", caller, timeout, objectName)
		}
	}()
}

// Destroy completely destroys the MrdInstance, cleaning up
// its resources and ensuring it is removed from the cache. This should be
// called when the owning inode is destroyed.
func (mi *MrdInstance) Destroy() {
	mi.refCountMu.Lock()
	defer mi.refCountMu.Unlock()

	// If it's in use, this indicates a potential lifecycle mismatch between the
	// inode and its readers.
	if mi.refCount > 0 {
		logger.Warnf("MrdInstance::Destroy called on an instance with refCount %d", mi.refCount)
	}

	// Remove from cache.
	if mi.mrdCache != nil {
		mi.mrdCache.Erase(getKey(mi.inodeId))
	}

	// Close the pool.
	mi.closePool()
}

// getKey generates a unique key for the MrdInstance based on its inode ID.
func getKey(id fuseops.InodeID) string {
	return strconv.FormatUint(uint64(id), 10)
}

// IncrementRefCount increases the reference count for the MrdInstance. When the
// instance is actively used (refCount > 0), it is removed from the inactive
// MRD cache to prevent eviction.
func (mi *MrdInstance) IncrementRefCount() {
	mi.refCountMu.Lock()
	defer mi.refCountMu.Unlock()
	mi.refCount++

	if mi.refCount == 1 && mi.mrdCache != nil {
		// Remove from cache
		deletedEntry := mi.mrdCache.Erase(getKey(mi.inodeId))
		if deletedEntry != nil {
			logger.Tracef("MrdInstance::IncrementRefCount: MrdInstance Inode (%d) erased from cache", mi.inodeId)
		}
	}
}

// handleEviction handles the cleanup of the MrdInstance when it is evicted from the cache.
// Race protection: MrdInstance could be reopened (refCount>0) or re-added to cache before eviction.
func (mi *MrdInstance) handleEviction() {
	mi.refCountMu.Lock()
	defer mi.refCountMu.Unlock()

	// Check if mrdInstance was reopened (refCount>0) - must skip eviction.
	if mi.refCount > 0 {
		return
	}

	// Check if mrdInstance was re-added to cache (refCount went 0→1→0 in between eviction and closure.)
	// Lock order: refCountMu -> cache.mu (consistent with Increment/DecrementRefCount)
	if mi.mrdCache != nil && mi.mrdCache.LookUpWithoutChangingOrder(getKey(mi.inodeId)) == mi {
		return
	}

	mi.closePool()
}

// DecrementRefCount decreases the reference count. When the count drops to zero, the
// instance is considered inactive and is added to the LRU cache for potential
// reuse. If the cache is full, this may trigger the eviction and closure of the
// least recently used MRD instances.
func (mi *MrdInstance) DecrementRefCount() {
	mi.refCountMu.Lock()
	defer mi.refCountMu.Unlock()

	if mi.refCount <= 0 {
		logger.Errorf("MrdInstance::DecrementRefCount: Refcount cannot be negative")
		return
	}

	mi.refCount--
	// Do nothing if MRDInstance is in use or cache is not enabled.
	if mi.refCount > 0 || mi.mrdCache == nil {
		return
	}

	// Add to cache.
	// Lock order: refCountMu -> cache.mu -> poolMu (via Size() inside Insert).
	// This is a safe order.
	evictedValues, err := mi.mrdCache.Insert(getKey(mi.inodeId), mi)
	if err != nil {
		logger.Errorf("MrdInstance::DecrementRefCount: Failed to insert MrdInstance for inode (%d) into cache, destroying immediately: %v", mi.inodeId, err)
		// The instance could not be inserted into the cache. Since the refCount is 0,
		// we must close it now to prevent it from being leaked.
		mi.closePool()
		return
	}

	logger.Tracef("MrdInstance::DecrementRefCount: MrdInstance for inode (%d) added to cache", mi.inodeId)

	// Do not proceed if no eviction happened.
	if evictedValues == nil {
		return
	}

	// Evict outside all locks.
	mi.refCountMu.Unlock()
	for _, instance := range evictedValues {
		mrdInstance, ok := instance.(*MrdInstance)
		if !ok {
			logger.Errorf("MrdInstance::DecrementRefCount: Invalid value type, expected *MrdInstance, got %T", instance)
		} else {
			mrdInstance.handleEviction()
		}
	}
	// Reacquire the lock ensuring safe defer's Unlock.
	mi.refCountMu.Lock()
}

// Size returns the number of active MRDs.
func (mi *MrdInstance) Size() uint64 {
	mi.poolMu.RLock()
	defer mi.poolMu.RUnlock()
	if mi.mrdPool != nil {
		return mi.mrdPool.Size()
	}
	return 0
}