Ayende @ Rahien

It's a girl

Answer: Stopping the leaks

Originally posted at 4/19/2011

Yesterday I posted the following challenge:

Given the following API, can you think of a way that would prevent memory leaks?

public interface IBufferPool
{
    byte[] TakeBuffer(int size);
    void ReturnBuffer(byte[] buffer);
}

The problem with having something like this is that forgetting to return the buffer is going to cause a memory leak. Instead of having that I would like to have the application stop if a buffer is leaked. Leaked means that no one is referencing this buffer but it wasn’t returned to the pool.

What I would really like is that when running in debug mode, leaking a buffer would stop the entire application and tell me:

  • That a buffer was leaked.
  • What was the stack trace that allocated that buffer.

Let us take a look at how we are going about implementing this, shall we? I am going to defer the actual implementation of the buffer pool to System.ServiceModel.Channels.BufferManager and focus on providing the anti leak features. The result is that this code:

IBufferPool pool = new BufferPool(1024*512, 1024);

var buffer = pool.TakeBuffer(512);
GC.WaitForPendingFinalizers(); // nothing here

pool.ReturnBuffer(buffer);
buffer = null;
GC.WaitForPendingFinalizers(); // nothing here, we released the memory properly

pool.TakeBuffer(512); // take and discard a buffer without returning to the pool
GC.WaitForPendingFinalizers(); // failure!

Will result in the following error:

Unhandled Exception: System.InvalidOperationException: A buffer was leaked. Initial allocation:
   at ConsoleApplication1.BufferPool.BufferTracker.TrackAllocation() in IBufferPool.cs:line 22
   at ConsoleApplication1.BufferPool.TakeBuffer(Int32 size) in IBufferPool.cs:line 60
   at ConsoleApplication1.Program.Main(String[] args) in Program.cs:line 21

And now for the implementation:

public class BufferPool : IBufferPool
{
    public class BufferTracker
    {
        private StackTrace stackTrace;

        public void TrackAllocation()
        {
            stackTrace = new StackTrace(true);
            GC.ReRegisterForFinalize(this);
        }

        public void Discard()
        {
            stackTrace = null;
            GC.SuppressFinalize(this);
        }

        ~BufferTracker()
        {
            if (stackTrace == null)
                return;

            throw new InvalidOperationException(
                "A buffer was leaked. Initial allocation:" + Environment.NewLine + stackTrace
                );
        }
    }

    private readonly BufferManager bufferManager;
    private ConditionalWeakTable<byte[], BufferTracker> trackLeakedBuffers = new ConditionalWeakTable<byte[], BufferTracker>();

    public BufferPool(long maxBufferPoolSize, int maxBufferSize)
    {
        bufferManager = BufferManager.CreateBufferManager(maxBufferPoolSize, maxBufferSize);
    }

    public void Dispose()
    {
        bufferManager.Clear();
        // note that disposing the pool before returning all of the buffers will cause a crash
    }

    public byte[] TakeBuffer(int size)
    {
        var buffer = bufferManager.TakeBuffer(size);
        trackLeakedBuffers.GetOrCreateValue(buffer).TrackAllocation();
        return buffer;
    }

    public void ReturnBuffer(byte[] buffer)
    {
        BufferTracker value;
        if(trackLeakedBuffers.TryGetValue(buffer, out value))
        {
            value.Discard();
        }
        bufferManager.ReturnBuffer(buffer);
    }
}

As you can see, utilizing ConditionalWeakTable is quite powerful, since it allows us to support a lot of really advanced scenarios in a fairly simple ways.

Comments

tobi
04/30/2011 11:54 AM by
tobi

Its most simple description is that the values have exactly the same lifetime as the keys. It is meant to attach additional data to existing, not extensible objects.

Ryan Heath
04/30/2011 12:23 PM by
Ryan Heath

How did I not notice ConditionalWeakTable before? I always felt a need for such behavior when extension methods were introduced to C# ...

Anyhow I did some googling and found someone had implemented my idea already, isnt that wonderful? :)

http://connectedproperties.codeplex.com/

// Ryan

Juan Lopes
04/30/2011 03:49 PM by
Juan Lopes

Are you sure you want to instantiate StackTrace for this task? Have you ever notice how slow is "new StackTrace()"?

Daniel Grunwald
04/30/2011 06:15 PM by
Daniel Grunwald

@Ryan: ConditionalWeakTable is new in .NET 4.0.

In fact, it is impossible to correctly implement the "connected properties" in .NET 3.5, since a special GC feature (ephemerons; new in .NET 4.0) is required to avoid leaking memory when the value references the key.

Unfortunately it seems that the WPF team at MS has't heard of ConditionalWeakTable yet - this memory leak in WPF still exists: http://support.microsoft.com/kb/938416/en-us

David
04/30/2011 08:49 PM by
David

This sort of approach reminds me of what Scott Bilas was calling Assertive Finalizers:

http://scottbilas.com/blog/assertive-finalizers/

Anyway, I'm surprised to see BufferPool have a Dispose method but not inherit the IDisposable interface - that would make it easier to use and have code analysis tools which verify that Dispose is called properly to work better. (Basically, when about to implement an interface like Dispose, I would recommend starting with typing the ": IDisposable", right-click and let it make the method for you, then fill in the code. As we just saw, the other way around can result in forgetting to add an interface which properly describes an important aspect of the class).

Erik
05/01/2011 12:29 PM by
Erik

I don't think your comment is entirely accurate: "// note that disposing the pool before returning all of the buffers will cause a crash"...

There is no such thing as an unhandled exception in the finalizer -- all that will happen is an error message will be output to the console. If that is true, how is this protecting us from leaked buffers?

I hadn't heard of ConditionalWeakTable before. That is amazing and extremely useful!

Ayende Rahien
05/01/2011 12:36 PM by
Ayende Rahien

Erik,

Unhandled exception in the finalizer will cause a crash, yes.

An error would be outputted to the console and the application will terminate.

Erik
05/01/2011 01:17 PM by
Erik

According to this: msdn.microsoft.com/en-us/library/ms228965.aspx

"There is no such thing as an unhandled exception on the finalizer thread. When a finalizer throws an exception that it does not handle, the runtime prints the exception stack trace to the console and then allows the finalizer thread to resume running finalizers."

It appears that this is a change in .NET4.

I could not reproduce appropriately failing conditions in a unit test either -- I could not get the expected exception thrown, even with GC.Collect's thrown in.

Ayende Rahien
05/01/2011 01:19 PM by
Ayende Rahien

Erik,

The documentation is wrong.

Try calling GC.WaitForPendingFinalizers();

Roja Buck
05/03/2011 08:54 AM by
Roja Buck

Ayende,

Forgive my ignorance but could you explain the reasoning / need for calling GC.ReRegisterForFinalize(this) within the BufferTracker constructor? As far as I can tell there is no way for SuppressFinalize to have been called on the object prior to it's construction and thus, as a finalizer is implemented in the class, would expect to find the newly created object on the finalization queue?
Ayende Rahien
05/03/2011 09:03 AM by
Ayende Rahien

Roja,

That isn't called in the constructor, BufferTracker doesn't HAVE a ctor.

Roja Buck
05/03/2011 09:15 AM by
Roja Buck

Oh dear,

Blinded by my lack of morning tea. What a buffoon!
Mistertom
05/10/2011 11:03 AM by
Mistertom

I always thought managed references were not supposed to be used in the finalizer.

RN
05/14/2011 05:03 PM by
RN

Erik, Ayende - the documentation above actually has Erik's quote under the section describing .NET framework 1.0 and 1.1 behaviour. In other words the quote describes the pre-Framework 2.0 behaviour.

Comments have been closed on this topic.