Skip to content

Commit

Permalink
Threas-safe mmap resize (#1891)
Browse files Browse the repository at this point in the history
  • Loading branch information
@BCSharp
BCSharp authored Jan 31, 2025
1 parent 990e5cd commit 3a477d2
Showing 1 changed file with 91 additions and 39 deletions.
130 changes: 91 additions & 39 deletions src/core/IronPython.Modules/mmap.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,20 +289,23 @@ public class MmapDefault : IWeakReferenceable {
private readonly MemoryMappedFileAccess _fileAccess;
private readonly SafeFileHandle? _handle;

// RefCount | Closed | Meaning
// ---------+--------+--------------------------------
// 0 | 0 | Not fully initialized
// 1 | 0 | Fully initialized, not being used by any threads
// >1 | 0 | Object in use by one or more threads
// >0 | 1 | Close/dispose requested, no more addrefs allowed, some threads may still be using it
// 0 | 1 | Fully disposed
private volatile int _state; // Combined ref count and closed/disposed flags (so we can atomically modify them).
// RefCount | Closed | Exclusive |Meaning
// ---------+--------+-----------+---------------------
// 0 | 0 | 0 | Not fully initialized
// 1 | 0 | 0 | Fully initialized, not being used by any threads
// >1 | 0 | 0 | Object in regular use by one or more threads
// 2 | 0 | 1 | Object in exclusive use (`resize` in progress)
// >0 | 1 | - | Close/dispose requested, no more addrefs allowed, some threads may still be using it
// 0 | 1 | 0 | Fully disposed
// Other combinations are invalid state
private volatile int _state; // Combined ref count and state flags (so we can atomically modify them).

private static class StateBits {
public const int Closed = 0b_01; // close/dispose requested; no more addrefs allowed
public const int Exclusive = 0b_10; // TODO: to manage exclusive access for resize
public const int RefCount = unchecked(~0b_11); // 2 bits reserved for state management; ref count gets 29 bits (sign bit unused)
public const int RefCountOne = 1 << 2; // ref count 1 shifted over 2 state bits
public const int Closed = 0b_001; // close/dispose requested; no more addrefs allowed
public const int Exclusive = 0b_010; // exclusive access for resize requested/in progress
public const int Exporting = 0b_100; // TODO: buffer exports extant; exclusive addrefs temporarily not allowed
public const int RefCount = unchecked(~0b_111); // 3 bits reserved for state management; ref count gets 29 bits (sign bit unused)
public const int RefCountOne = 1 << 3; // ref count 1 shifted over 3 state bits
}


Expand Down Expand Up @@ -563,33 +566,74 @@ public void __exit__(CodeContext/*!*/ context, params object[] excinfo) {
public bool closed => (_state & StateBits.Closed) == StateBits.Closed; // Dispose already requested, will self-dispose when ref count drops to 0.


private bool AddRef() {
/// <summary>
/// Try to add a reference to the mmap object. Return <c>true</c> on success.
/// </summary>
/// <remarks>
/// The reference count is incremented atomically and kept in the mmap state variable.
/// The reference count is not incremented if the state variable indicates that the mmap
/// in in the process of closing or currently in exclusive use.
/// </remarks>
/// <param name="exclusive">
/// If true, requests an exclusive reference.
/// </param>
/// <param name="reason">
/// If the reference could not be added, this parameter will contain the bit that was set preventing the addref.
/// </param>
private bool TryAddRef(bool exclusive, out int reason) {
int oldState, newState;
do {
oldState = _state;
if ((oldState & StateBits.Closed) == StateBits.Closed) {
// mmap closed, dispose already requested, no more addrefs allowed
reason = StateBits.Closed;
return false;
}
if ((oldState & StateBits.Exclusive) == StateBits.Exclusive) {
// mmap in exclusive use, temporarily no more addrefs allowed
reason = StateBits.Exclusive;
return false;
}
if (exclusive && (oldState & StateBits.Exporting) == StateBits.Exporting) {
// mmap exporting, exclusive addrefs temporarily not allowed
reason = StateBits.Exporting;
return false;
}
Debug.Assert((oldState & StateBits.RefCount) > 0, "resurrecting disposed mmap object (disposed without being closed)");

newState = oldState + StateBits.RefCountOne;
if (exclusive) {
newState |= StateBits.Exclusive;
}
} while (Interlocked.CompareExchange(ref _state, newState, oldState) != oldState);
reason = 0;
return true;
}


private void Release() {
/// <summary>
/// Atomically release a reference to the mmap object, and optionally reset the exclusive state flag.
/// </summary>
/// <remarks>
/// If the reference count drops to 0, the mmap object is disposed.
/// </remarks>
/// <param name="exclusive">
/// If true, the exclusive reference is released.
/// </param>
private void Release(bool exclusive) {
bool performDispose;
int oldState, newState;
do {
oldState = _state;
Debug.Assert(!exclusive || (oldState & StateBits.Exclusive) == StateBits.Exclusive, "releasing exclusive reference without being exclusive");
Debug.Assert((oldState & StateBits.RefCount) > 0, "mmap ref count underflow (too many releases)");

performDispose = (oldState & StateBits.RefCount) == StateBits.RefCountOne;
Debug.Assert(!performDispose || (oldState & StateBits.Closed) == StateBits.Closed, "disposing mmap object without being closed");

newState = oldState - StateBits.RefCountOne;
if (performDispose) {
newState |= StateBits.Closed; // most likely already closed
if (exclusive) {
newState &= ~StateBits.Exclusive;
}
} while (Interlocked.CompareExchange(ref _state, newState, oldState) != oldState);

Expand All @@ -610,25 +654,17 @@ public void close() {
#if NET5_0_OR_GREATER
if ((Interlocked.Or(ref _state, StateBits.Closed) & StateBits.Closed) != StateBits.Closed) {
// freshly closed, release the construction time reference
Release();
Release(exclusive: false);
}
#else
int current = _state;
while (true)
{
int newState = current | StateBits.Closed;
int oldState = Interlocked.CompareExchange(ref _state, newState, current);
if (oldState == current)
{
// didn't change in the meantime, exchange with newState completed
if ((oldState & StateBits.Closed) != StateBits.Closed) {
// freshly closed, release the construction time reference
Release();
}
return;
}
// try again to set the bit
current = oldState;
int oldState, newState;
do {
oldState = _state;
newState = oldState | StateBits.Closed;
} while (Interlocked.CompareExchange(ref _state, newState, oldState) != oldState);
if ((oldState & StateBits.Closed) != StateBits.Closed) {
// freshly closed, release the construction time reference
Release(exclusive: false);
}
#endif
}
Expand Down Expand Up @@ -861,8 +897,9 @@ public string readline() {
}
}


public void resize(long newsize) {
using (new MmapLocker(this)) {
using (new MmapLocker(this, exclusive: true)) {
if (_fileAccess is not MemoryMappedFileAccess.ReadWrite and not MemoryMappedFileAccess.ReadWriteExecute) {
throw PythonOps.TypeError("mmap can't resize a readonly or copy-on-write memory map.");
}
Expand Down Expand Up @@ -961,6 +998,7 @@ public void resize(long newsize) {
}
}


public object rfind([NotNone] IBufferProtocol s) {
using (new MmapLocker(this)) {
return RFindWorker(s, Position, _view.Capacity);
Expand Down Expand Up @@ -1210,18 +1248,32 @@ private static long GetFileSizeUnix(SafeFileHandle handle) {

private readonly struct MmapLocker : IDisposable {
private readonly MmapDefault _mmap;

public MmapLocker(MmapDefault mmap) {
if (!mmap.AddRef()) {
throw PythonOps.ValueError("mmap closed or invalid");
private readonly bool _exclusive;

public MmapLocker(MmapDefault mmap, bool exclusive = false) {
if (!mmap.TryAddRef(exclusive, out int reason)) {
if (reason == StateBits.Closed) {
// mmap is permanently closed
throw PythonOps.ValueError("mmap closed or invalid");
} else if (reason == StateBits.Exporting) {
// map is temporarily exporting buffers obtained through the buffer protocol
throw PythonOps.BufferError("mmap can't perform the operation with extant buffers exported");
} else if (reason == StateBits.Exclusive) {
// mmap is temporarily in exclusive use
throw PythonNT.GetOsError(PythonErrno.EAGAIN);
} else {
// should not happen
throw new InvalidOperationException("mmap state error");
}
}
_mmap = mmap;
_exclusive = exclusive;
}

#region IDisposable Members

public readonly void Dispose() {
_mmap.Release();
_mmap.Release(_exclusive);
}

#endregion
Expand Down

0 comments on commit 3a477d2

Please sign in to comment.