ProjectArcade

1
using System;
2
using System.Collections.Concurrent;
3
using System.Threading;
4
using DokanNet.Logging;
5

6
namespace DokanNet
7
{
8
    /// <summary>
9
    /// Simple buffer pool for buffers used by <see cref="IDokanOperations.ReadFile"/> and
10
    /// <see cref="IDokanOperations.WriteFile"/> to avoid excessive Gen2 garbage collections due to large buffer
11
    /// allocation on the large object heap (LOH).
12
    /// 
13
    /// This pool is a bit different than say System.Buffers.ArrayPool(T) in that <see cref="RentBuffer" /> only returns
14
    /// exact size buffers. This is because the Read/WriteFile APIs only take a byte[] array as a parameter, not
15
    /// buffer and length. As such, it would be back compat breaking to return buffers that are bigger than the
16
    /// data length. To limit the amount of memory consumption, we only buffer sizes that are powers of 2 because
17
    /// common buffer sizes are typically that. There isn't anything preventing pooling buffers of any size if
18
    /// we find that there's another common buffer size in use. Only pool buffers 1MB or smaller and only
19
    /// up to 10 buffers of each size for further memory capping.
20
    /// </summary>
21
    internal class BufferPool
22
    {
23
        // An empty array does not contain data and can be statically cached.
24
        private static readonly byte[] _emptyArray = new byte[0];
25

26
        private readonly uint _maxBuffersPerPool; // Max buffers to cache per buffer size.
27

28
        // The pools for each buffer size. Index is log2(size).
29
        private readonly ConcurrentBag<byte[]>[] _pools;
30

31
        // Number of bytes served out over the pool's lifetime.
32
        private long _servedBytes;
33

34
        /// <summary>
35
        /// Constructs a new buffer pool.
36
        /// </summary>
37
        /// <param name="maxBufferSize">The max size (bytes) buffer that will be cached. </param>
38
        /// <param name="maxBuffersPerBufferSize">Maximum number of buffers cached per buffer size.</param>
39
        public BufferPool(uint maxBufferSize = 1024 * 1024, uint maxBuffersPerBufferSize = 10)
40
        {
41
            _maxBuffersPerPool = maxBuffersPerBufferSize;
42
            int log2 = GetPoolIndex(maxBufferSize);
43
            if (log2 == -1)
44
            {
45
                throw new ArgumentOutOfRangeException("maxBufferSize", maxBufferSize, "Must be a power of 2.");
46
            }
47

48
            // Create empty pools for each size.
49
            _pools = new ConcurrentBag<byte[]>[log2 + 1];
50
            for (int i = 0; i < _pools.Length; i++)
51
            {
52
                _pools[i] = new ConcurrentBag<byte[]>();
53
            }
54
        }
55

56
        /// <summary>
57
        /// Default, process-wide buffer pool instance.
58
        /// </summary>
59
        public static BufferPool Default { get { return _default; } }
60
        private static BufferPool _default = new BufferPool();
61

62
        /// <summary>
63
        /// Clears the buffer pool by releasing all buffers.
64
        /// </summary>
65
        public void Clear()
66
        {
67
            _servedBytes = 0;
68
            for (int i = 0; i < _pools.Length; i++)
69
            {
70
                _pools[i] = new ConcurrentBag<byte[]>(); // There's no clear method on ConcurrentBag...
71
            }
72
        }
73

74
        /// <summary>
75
        /// Number of bytes served over the pool's lifetime.
76
        /// </summary>
77
        public long ServedBytes { get { return Interlocked.Read(ref _servedBytes); } }
78

79
        /// <summary>
80
        /// Gets a buffer from the buffer pool of the exact specified size.
81
        /// If the size if not a power of 2, a buffer is still returned, but it is not poolable.
82
        /// </summary>
83
        /// <param name="bufferSize">The size of buffer requested.</param>
84
        /// <param name="logger">Logger for debug spew about what the buffer pool did.</param>
85
        /// <returns>The byte[] buffer.</returns>
86
        public byte[] RentBuffer(uint bufferSize, ILogger logger)
87
        {
88
            if (bufferSize == 0)
89
            {
90
                return _emptyArray; // byte[0] is statically cached.
91
            }
92

93
            Interlocked.Add(ref _servedBytes, bufferSize);
94

95
            // If the number is not a power of 2, we have nothing to offer.
96
            int poolIndex = GetPoolIndex(bufferSize);
97
            if (poolIndex == -1 || poolIndex >= _pools.Length)
98
            {
99
                if (logger.DebugEnabled) logger.Debug("Buffer size "+bufferSize+" not power of 2 or too large, returning unpooled buffer.");
100
                return new byte[bufferSize];
101
            }
102

103
            // Try getting a buffer from the pool. If it's empty, make a new buffer.
104
            ConcurrentBag<byte[]> pool = _pools[poolIndex];
105
            byte[] buffer;
106
            if (pool.TryTake(out buffer))
107
            {
108
                if (logger.DebugEnabled) logger.Debug("Using pooled buffer from pool "+poolIndex);
109
            }
110
            else
111
            {
112
                if (logger.DebugEnabled) logger.Debug("Pool "+poolIndex+" empty, creating new buffer.");
113
                buffer = new byte[bufferSize];
114
            }
115

116
            return buffer;
117
        }
118

119
        /// <summary>
120
        /// Returns a previously rented buffer to the buffer pool.
121
        /// If the buffer size is not an exact power of 2, the buffer is ignored.
122
        /// </summary>
123
        /// <param name="buffer">The buffer to return.</param>
124
        /// <param name="logger">Logger for debug spew about what the buffer pool did.</param>
125
        public void ReturnBuffer(byte[] buffer, ILogger logger)
126
        {
127
            if (buffer.Length == 0)
128
            {
129
                return; // Do nothing - _emptyArray caches this statically.
130
            }
131

132
            // If the buffer is a power of 2 and below max pooled size, return it to the appropriate pool.
133
            int poolIndex = GetPoolIndex((uint)buffer.Length);
134
            if (poolIndex >= 0 && poolIndex < _pools.Length)
135
            {
136
                // Check if the pool is full. This is racy if multiple threads return buffers concurrently,
137
                // but it's close enough - we'd just get a couple extra buffers in the pool at worst.
138
                ConcurrentBag<byte[]> pool = _pools[poolIndex];
139
                if (pool.Count < _maxBuffersPerPool)
140
                {
141
                    Array.Clear(buffer, 0, buffer.Length);
142
                    pool.Add(buffer);
143
                    if (logger.DebugEnabled) logger.Debug("Returned buffer to pool " +poolIndex);
144
                }
145
                else
146
                {
147
                    if (logger.DebugEnabled) logger.Debug("Pool " + poolIndex + " is full, discarding buffer.");
148
                }
149
            }
150
            else
151
            {
152
                if (logger.DebugEnabled) logger.Debug(poolIndex + " size "+ buffer.Length + ") outside pool range, discarding buffer.");
153
            }
154
        }
155

156
        /// <summary>
157
        /// Computes the pool index given a buffer size. The pool index is log2(size),
158
        /// if size is a power of 2. If size is not a power of 2, -1 is returned (invalid pool index).
159
        /// </summary>
160
        /// <param name="bufferSize">Buffer size in bytes.</param>
161
        /// <returns>The pool index, log2(number), or -1 if bufferSize is not a power of 2.</returns>
162
        private static int GetPoolIndex(uint bufferSize)
163
        {
164
            double log2 = Math.Log(bufferSize, 2);
165
            int log2AsInt = (int)log2;
166

167
            // If they are not equal, the number is not a power of 2.
168
            //
169
            if (log2 != log2AsInt)
170
            {
171
                return -1;
172
            }
173

174
            return log2AsInt;
175
        }
176
    }
177
}
178