CloudFetchMemoryBufferManager.cs

/*
* Copyright (c) 2025 ADBC Drivers Contributors
*
* This file has been modified from its original version, which is
* under the Apache License:
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you 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.
*/

using System;
using System.Threading;
using System.Threading.Tasks;

namespace AdbcDrivers.Databricks.Reader.CloudFetch
{
    /// <summary>
    /// Manages memory allocation for prefetched files.
    ///
    /// This manager tracks in-flight download memory based on the compressed file sizes
    /// received from the server. It does NOT track the decompressed data size, as that
    /// memory is naturally bounded by the result queue capacity and batch processing flow.
    ///
    /// The memory limit controls how many concurrent downloads can be in-flight at once,
    /// preventing unbounded parallel downloads from exhausting system resources.
    /// </summary>
    internal sealed class CloudFetchMemoryBufferManager : ICloudFetchMemoryBufferManager
    {
        private readonly long _maxMemory;
        private long _usedMemory;
        private readonly SemaphoreSlim _memorySemaphore;

        /// <summary>
        /// Initializes a new instance of the <see cref="CloudFetchMemoryBufferManager"/> class.
        /// </summary>
        /// <param name="maxMemoryMB">The maximum memory allowed for buffering in megabytes.</param>
        public CloudFetchMemoryBufferManager(int maxMemoryMB = CloudFetchConfiguration.DefaultMemoryBufferSizeMB)
        {
            if (maxMemoryMB <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(maxMemoryMB), "Memory buffer size must be positive.");
            }

            // Convert MB to bytes
            _maxMemory = maxMemoryMB * 1024L * 1024L;
            _usedMemory = 0;
            _memorySemaphore = new SemaphoreSlim(1, 1);
        }

        /// <inheritdoc />
        public long MaxMemory => _maxMemory;

        /// <inheritdoc />
        public long UsedMemory => Interlocked.Read(ref _usedMemory);

        /// <inheritdoc />
        public bool TryAcquireMemory(long size)
        {
            if (size <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(size), "Size must be positive.");
            }

            // Try to acquire memory (size is the compressed file size from the server)
            long originalValue;
            long newValue;
            do
            {
                originalValue = Interlocked.Read(ref _usedMemory);
                newValue = originalValue + size;

                // Check if we would exceed the maximum memory
                if (newValue > _maxMemory)
                {
                    return false;
                }
            }
            while (Interlocked.CompareExchange(ref _usedMemory, newValue, originalValue) != originalValue);

            return true;
        }

        /// <inheritdoc />
        public async Task AcquireMemoryAsync(long size, CancellationToken cancellationToken)
        {
            if (size <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(size), "Size must be positive.");
            }

            // Special case: if size is greater than max memory, we'll never be able to acquire it
            // Note: size is the compressed file size from the server
            if (size > _maxMemory)
            {
                throw new ArgumentOutOfRangeException(nameof(size), $"Requested size ({size} bytes) exceeds maximum memory ({_maxMemory} bytes).");
            }

            while (!cancellationToken.IsCancellationRequested)
            {
                // Try to acquire memory without blocking
                if (TryAcquireMemory(size))
                {
                    return;
                }

                // If we couldn't acquire memory, wait for some to be released
                await Task.Delay(10, cancellationToken).ConfigureAwait(false);
            }

            // If we get here, cancellation was requested
            cancellationToken.ThrowIfCancellationRequested();
        }

        /// <inheritdoc />
        public void ReleaseMemory(long size)
        {
            if (size <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(size), "Size must be positive.");
            }

            // Release memory
            long newValue = Interlocked.Add(ref _usedMemory, -size);

            // Ensure we don't go negative
            if (newValue < 0)
            {
                // This should never happen if the code is correct
                Interlocked.Exchange(ref _usedMemory, 0);
                throw new InvalidOperationException("Memory buffer manager released more memory than was acquired.");
            }
        }
    }
}