3
\$\begingroup\$

I've recently needed a class to help me manage a pool of re-usable objects. It's my first time delving into the async/await side of C# and I wanted to create an async-friendly interface to allow me to elegantly handle situations where no pool item is currently available. I've come up with the below solution, but I'm not sure it's the best way to implement such a system.

The main concerns I have are with when all of the pool items are exhausted and calls to Acquire need to wait until an item is available. The current implementation needs to create a new TaskCompletionSource for each new call to Acquire once the pool is exhausted.

The usage pattern could be something like this:

var pool = new AsyncObjectPool<T>(...);

while(true) {
    // Wait for an object to become available
    var obj = await pool.Acquire(CancellationToken.None);
    // Example method, wait for an incoming request
    await obj.ReceiveRequestAsync();
    // Process the received request, free'ing the object at the end of it
    // The loop immediately continues to wait for the next incoming request
    ProcessRequest(obj).ContinueWith(_ => pool.Free(obj));
}

In this situation the pool fills up to capacity fairly quickly, then acts as a gatekeeper to prevent too many requests from processing concurrently as there aren't enough resources. So I suppose my main questions are:

  • Is this code correct and thread-safe (i.e. no obvious deadlocks, race conditions, etc.)
  • Is this even the right pattern to take for the above use-case? The main reason standard queuing wasn't an option was that I actually needed to wait for the object to become available before being able to check/wait for an incoming request.

The code:

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

namespace AsyncTools {
    public class AsyncObjectPool<T> where T : class {
        private const int OBJECT_BORROWED = 1;
        private const int OBJECT_AVAILABLE = 0;

        private readonly PoolItem[] _pool;
        private readonly Func<int, T> _factory;
        private readonly LinkedList<TaskCompletionSource<T>> _waiting;

        public AsyncObjectPool(int capacity, Func<int, T> factory) {
            _pool = new PoolItem[capacity];
            _waiting = new LinkedList<TaskCompletionSource<T>>();
            _factory = factory;
        }

        public async Task<T> Acquire(CancellationToken token) {
            var localPool = _pool;

            for(int i = 0; i < localPool.Length; ++i) {
                if (Interlocked.CompareExchange(ref localPool[i].State, OBJECT_BORROWED, OBJECT_AVAILABLE) == OBJECT_AVAILABLE) {
                    return localPool[i].Value ?? (localPool[i].Value = _factory(i));
                }
            }

            var tcs = new TaskCompletionSource<T>();

            lock(_waiting) 
                _waiting.AddLast(tcs);

            if (!token.CanBeCanceled)
                return await tcs.Task.ConfigureAwait(false);

            // Configure cancellation handling
            var cancellationHandler = new Action(() => {
                if (!tcs.Task.IsCompleted) {
                    lock(_waiting)
                        _waiting.Remove(tcs);
                    tcs.SetCanceled();
                }
            });

            using(token.Register(cancellationHandler))
                return await tcs.Task.ConfigureAwait(false);
        }

        public void Free(T obj) {
            if (obj == null)
                throw new ArgumentNullException(nameof(obj));

            // Deliberately avoid taking a lock on _waiting pre-emptively. Worst case
            // scenario is a recently added "waiter" gets skipped until the next
            // resource is free'd which is extremely unlikely.
            if (_waiting.Count == 0) {
                for(int i = 0; i < _pool.Length; ++i) {
                    // Each object should only be owned by one user at a time,
                    // which means this read shouldn't need to be interlocked
                    if (_pool[i].Value == obj) {
                        Interlocked.Exchange(ref _pool[i].State, OBJECT_AVAILABLE);
                        break;
                    }
                }
            }
            else {
                TaskCompletionSource<T> tcs = null;

                lock(_waiting) {
                    // Double-read 'Count' as we may have lost a contentious lock
                    // with only a single item in the list since our last read.
                    if (_waiting.Count > 0) {
                        tcs = _waiting.First.Value;
                        _waiting.RemoveFirst();
                    }
                }

                if (tcs != null)
                    tcs.SetResult(obj);
            }
        }

        public void Remove(T obj) {
            if (obj == null)
                throw new ArgumentNullException(nameof(obj));

            // Same semantics as 'Free' in that only one PoolItem can own this resource
            // at a time so there should be no need for interlocked reads
            int i = 0;
            for (; i < _pool.Length; ++i) {
                if (_pool[i].Value == obj)
                    break;
            }

             // Was the item even in the pool?
            if (i >= _pool.Length)
                return;

            TaskCompletionSource<T> tcs = null;

            // If there are any pending 'waiters' we can allocate one a new instance
            // and resolve it right away. 

            // PERF: Can remove this if you don't care too much about the order that 
            //  waiters are resolved. Without this, the PoolItem will be claimed by the
            //  next call to 'Acquire' and the waiter will have to wait for a resource from
            //  'Free' instead. There is the potential for a deadlock if there are waiters,
            //  no further calls to Acquire occur, and every single current object is removed rather than freed.
            if (_waiting.Count > 0) {
                lock(_waiting) {
                    // Double-read 'Count' as we may have lost a contentious lock
                    // with only a single item in the list since our last read.
                    if (_waiting.Count > 0) {
                        tcs = _waiting.First.Value;
                        _waiting.RemoveFirst();
                    }
                }
            }

            if (tcs != null) 
                tcs.SetResult(_pool[i].Value = _factory(i));
            else {
                // Noone needs an object right now, just free the item slot
                _pool[i].Value = null;
                Interlocked.Exchange(ref _pool[i].State, OBJECT_AVAILABLE);
            }
        }

        private struct PoolItem {
            public T Value;
            public int State;
        }
    }
}
\$\endgroup\$
2
\$\begingroup\$

The first thing to strike me as a definite, yet subtle, bug:

                    if (_pool[i].Value == obj) {

This is fine until someone uses the class with a T which overrides ==. In an object pool you care about object identity, so this should use object.ReferenceEquals. (NB the above line occurs twice).


The issue of concurrency correctness is tricky. I understand your preference for avoiding pre-emptive locking, but it is by far the easiest way to get things right.

            // Deliberately avoid taking a lock on _waiting pre-emptively. Worst case
            // scenario is a recently added "waiter" gets skipped until the next
            // resource is free'd which is extremely unlikely.

is wrong. The worst case scenario is that one thread is in the Acquire method, about to acquire the lock, and while it's in that state every single object gets freed. That's a deadlock: given the example use case, the loop is blocked unless there's another thread which can make a parallel call to Acquire followed by Free. Moreover, expanding the scope of the lock in Free and Remove isn't sufficient: the scope of the lock in Acquire is also too narrow.


There's also a trade-off between complexity and trusting your user. The example use case calls Free, but if someone uses the pool badly they could leak the objects by losing references to them without calling Free. If that only occurs in a rare error case it could pass unnoticed for a long time. The paranoid option would be to return an IDisposable wrapper around T which also has a destructor to ensure that if the reference is lost it automatically frees the wrapped object.

\$\endgroup\$
  • \$\begingroup\$ Thanks for the follow-up, I had considered the IDisposable pattern, however I was hoping to reduce the number of allocations per Acquire call (which is why creating a new TokenCompletionSource for each waiting request was a pain). Probably a massive premature optimization though. I reasoned that this is a class for my own use, and I can accept the responsibility for ensuring objects are Free'd back to the pool. If I mess it up I only have my self to blame. I'd probably do it differently if it was a public library. \$\endgroup\$ – Jason Larke Sep 8 '17 at 1:29
  • \$\begingroup\$ Also thanks for pointing out the deadlock, even re-reading my own comment makes it clear that there's a potential deadlock condition I just hadn't followed it far enough to realise, d'oh \$\endgroup\$ – Jason Larke Sep 8 '17 at 1:47

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.