DRY RAII with AutoBuffer

Don’t Repeat Yourself is a good maxim in software design — avoid duplication of information. Today’s sample code will show a small fix to a problem which quite literally involves repeating yourself: those annoying “call twice” Win32 functions like GetVirtualDiskPhysicalPath.

Here’s an example of how one might call it in “old school” C-style code:

ULONG size = 16;
WCHAR * path = new WCHAR[size];
DWORD error = GetVirtualDiskPhysicalPath(handle, &size, path);
if (error == ERROR_INSUFFICIENT_BUFFER)
{
    // Sigh. Call again now that we know the real size to use...
    delete [] path;
    path = new WCHAR[size];
    error = GetVirtualDiskPhysicalPath(handle, &size, path)
}

// ... more code here, which will hopefully eventually free the buffer ...

There are multiple problems here. The duplication is ugly. The naked memory management calls leave you prone to leaks if you don’t remember to free on every return path. This “call twice to get buffer size” is also a common pattern for many other Win32 calls, so what you do here has to be structurally repeated in many places.

With C++11 (or better), you can elegantly solve this problem with judicious use of lambda expressions, STL vectors, and templating. Here is my stab at it in a class called AutoBuffer. It attempts to provide a type-safe and RAII-observant solution to the problem above.

Let’s rewrite the previous sample code using AutoBuffer:

AutoBuffer<WCHAR> buffer;
DWORD error = buffer.Fill<ULONG>([handle](WCHAR * dataPtr, PULONG sizePtr)
{
    return GetVirtualDiskPhysicalPath(handle, sizePtr, dataPtr);
});

// more code here... but no need to free the buffer, since AutoBuffer will destroy it for you!

Much simpler, eh? No more duplication and a lower chance of bugs.

Below is a reference implementation with unit tests that takes care of the “return error code” variation of this pattern, as demonstrated above. (Note that there are other Win32 functions using the “call twice” semantics such as QueryServiceConfig that instead return BOOL and GetLastError must be used; I’ll leave that one as an exercise for the reader….)

// Implementation -- AutoBuffer.h
#pragma once

#include <Windows.h>
#include <vector>
#include <functional>

template<typename TData>
class AutoBuffer
{
public:
    AutoBuffer()
        : buffer_(sizeof(TData), '\0')
    {
    }

    ~AutoBuffer()
    {
    }

    template<typename TSize>
    DWORD Fill(std::function<DWORD(TData *, TSize *)> func)
    {
        TSize size = static_cast<TSize>(get_Size());
        DWORD result = func(Ptr(0), &size);
        if (result == ERROR_INSUFFICIENT_BUFFER)
        {
            buffer_.resize(size);
            result = func(Ptr(0), &size);
        }

        return result;
    }

    size_t get_Size() const
    {
        return buffer_.size();
    }

    TData * operator->()
    {
        return Ptr(0);
    }

    TData & operator[](int index)
    {
        return *Ptr(index);
    }

private:
    std::vector<unsigned char> buffer_;

    TData * Ptr(int index)
    {
        return reinterpret_cast<TData *>(&buffer_[index * sizeof(TData)]);
    }
};

// Tests -- AutoBufferTest.cpp
#include "CppUnitTest.h"
#include "AutoBuffer.h"

using namespace Microsoft::VisualStudio::CppUnitTestFramework;

namespace UnitTest1
{
    TEST_CLASS(AutoBufferTest)
    {
    public:
        struct ExampleData
        {
            int Number;
            void * Ptr;
        };


        TEST_METHOD(ShouldInitializeWithSizeOfOneItemEmpty)
        {
            AutoBuffer<ExampleData> buffer;

            Assert::AreEqual(sizeof(ExampleData), buffer.get_Size());
            Assert::AreEqual(0, buffer->Number);
            Assert::IsNull(buffer->Ptr);
        }

        TEST_METHOD(ShouldFillAndReturnSuccessWhenFirstCallSucceeds)
        {
            AutoBuffer<ExampleData> buffer;

            DWORD result = buffer.Fill<ULONG>([](ExampleData * dataPtr, PULONG sizePtr)
            {
                dataPtr->Number = *sizePtr;
                return ERROR_SUCCESS;
            });

            Assert::AreEqual(static_cast<DWORD>(0), result);
            Assert::AreEqual(static_cast<int>(sizeof(ExampleData)), buffer->Number);
        }

        TEST_METHOD(ShouldReturnInitialErrorWhenFirstCallFails)
        {
            AutoBuffer<ExampleData> buffer;

            DWORD result = buffer.Fill<ULONG>([](ExampleData * dataPtr, PULONG sizePtr)
            {
                return ERROR_NOT_FOUND;
            });

            Assert::AreEqual(static_cast<DWORD>(ERROR_NOT_FOUND), result);
            Assert::AreEqual(0, buffer->Number);
        }

        TEST_METHOD(ShouldCallAgainWithUpdatedSizeIfFirstCallFailsWithInsufficientBuffer)
        {
            AutoBuffer<ExampleData> buffer;

            DWORD result = buffer.Fill<ULONG>([](ExampleData * dataPtr, PULONG sizePtr)
            {
                if (++dataPtr->Number == 1)
                {
                    *sizePtr *= 2;
                    return ERROR_INSUFFICIENT_BUFFER;
                }

                dataPtr[1].Number = 1000;
                return ERROR_SUCCESS;
            });

            Assert::AreEqual(static_cast<DWORD>(0), result);
            Assert::AreEqual(2, buffer->Number);
            Assert::AreEqual(1000, buffer[1].Number);
        }

        TEST_METHOD(ShouldReturnFinalErrorIfSecondCallFails)
        {
            AutoBuffer<ExampleData> buffer;

            DWORD result = buffer.Fill<ULONG>([](ExampleData * dataPtr, PULONG sizePtr)
            {
                if (++dataPtr->Number == 1)
                {
                    return ERROR_INSUFFICIENT_BUFFER;
                }

                return ERROR_NOT_FOUND;
            });

            Assert::AreEqual(static_cast<DWORD>(ERROR_NOT_FOUND), result);
        }
    };
}