StripedWorkerQueue.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.Threading;

namespace kinetica;

    internal sealed class Stripe<T>
    {
        private readonly object _lock = new();
        private List<T> _records;
        private readonly int _capacity;

        public Stripe(int capacity)
        {
            _capacity = capacity;
            _records = new List<T>(capacity);
        }

        public List<T>? Add(T record)
        {
            lock (_lock)
            {
                _records.Add(record);
                if (_records.Count >= _capacity)
                {
                    var batch = _records;
                    _records = new List<T>(_capacity);
                    return batch;
                }
                return null;
            }
        }

        public List<List<T>> AddRange(IReadOnlyList<T> records)
        {
            var batches = new List<List<T>>();
            lock (_lock)
            {
                foreach (var record in records)
                {
                    _records.Add(record);
                    if (_records.Count >= _capacity)
                    {
                        batches.Add(_records);
                        _records = new List<T>(_capacity);
                    }
                }
            }
            return batches;
        }

        public List<T>? Flush()
        {
            lock (_lock)
            {
                if (_records.Count == 0)
                    return null;

                var batch = _records;
                _records = new List<T>(_capacity);
                return batch;
            }
        }

        public int Count
        {
            get
            {
                lock (_lock)
                {
                    return _records.Count;
                }
            }
        }
    }

    internal sealed class StripedWorkerQueue<T>
    {
        private readonly Stripe<T>[] _stripes;
        private readonly int _numStripes;
        private readonly int _stripeMask;

        public Uri Url { get; }

        public int WorkerIndex { get; }

        public StripedWorkerQueue(Uri url, int workerIndex, int numStripes, int batchSize)
        {
            Url = url;
            WorkerIndex = workerIndex;

            // Round up to nearest power of 2 for efficient modulo operation
            _numStripes = RoundUpToPowerOf2(numStripes);
            _stripeMask = _numStripes - 1;

            _stripes = new Stripe<T>[_numStripes];
            for (int i = 0; i < _numStripes; i++)
            {
                _stripes[i] = new Stripe<T>(batchSize);
            }
        }

        private static int RoundUpToPowerOf2(int n)
        {
            n--;
            n |= n >> 1;
            n |= n >> 2;
            n |= n >> 4;
            n |= n >> 8;
            n |= n >> 16;
            n++;
            return Math.Max(1, n);
        }

        public List<T>? Add(T record, long stripeHash, out int stripeIndex)
        {
            stripeIndex = (int)(stripeHash & _stripeMask);
            return _stripes[stripeIndex].Add(record);
        }

        public List<T>? AddToStripe(T record, int stripeIndex)
        {
            return _stripes[stripeIndex & _stripeMask].Add(record);
        }

        public List<List<T>> AddRangeToStripe(IReadOnlyList<T> records, int stripeIndex)
        {
            return _stripes[stripeIndex & _stripeMask].AddRange(records);
        }

        public List<(int stripeIndex, List<T> batch)> FlushAll()
        {
            var batches = new List<(int, List<T>)>();
            for (int i = 0; i < _numStripes; i++)
            {
                var batch = _stripes[i].Flush();
                if (batch != null && batch.Count > 0)
                {
                    batches.Add((i, batch));
                }
            }
            return batches;
        }

        public int TotalCount
        {
            get
            {
                int count = 0;
                for (int i = 0; i < _numStripes; i++)
                {
                    count += _stripes[i].Count;
                }
                return count;
            }
        }

        public int NumStripes => _numStripes;
    }