Answer: Stopping the leaks

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.