p_tan's blog

勉強日記です。ツッコミ大歓迎

CoTaskMemAllocのリークをIMallocSpy で検知する

MSVCでメモリリークチェックと言えば、CRT ライブラリを使用したメモリ リークの検出 が定番ですが、これだとCランタイム経由じゃない CoTaskMemAlloc, CoTaskMemReallocで確保したメモリのリークは検知できません。 そこで IMallocSpy interface (COM) を使います。

IMallocSpy を使用したリークチェック

CoTaskMemAlloc は結局 IMalloc::Allocを呼びに行くのですが、IMallocSpy を CoRegisterMallocSpy() 関数でプロセスに登録しておくと、IMallocの各種メソッドが呼ばれる前後にIMallocSpyの各メソッドが呼ばれるようになるので、そこに好きなコードを書けばリークチェックとかが出来るようになります。 コードは以下。

#include <ObjIdl.h>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <unordered_map>

using namespace std;

// リークチェック用
class MallocSpy : public IMallocSpy
{
    // メモリ割当情報.
    struct AllocInfo {
        size_t count = 0;     // 割当番号
        SIZE_T allocSize = 0;  // 割当サイズ

        void set(SIZE_T size) {
            ++count;
            allocSize = size;
        }
    } m_allocInfo;
    using AllocMap = unordered_map<void*, AllocInfo>;
    AllocMap m_alloced; // アドレスから割当情報へのマップ

public:
    bool NoLeak() const {
        return m_alloced.empty();
    }
    // リーク情報を文字列にダンプする
    string DumpLeak() const {
        if (NoLeak()) { return ""s; }
        stringstream ss;
        vector<pair<void*, AllocInfo>> leaks(m_alloced.begin(), m_alloced.end());
        std::sort(leaks.begin(), leaks.end(), [](const auto &lhs, const auto &rhs) {
            return lhs.second.count < rhs.second.count;
        });
        ss << "CoTaskMem Leak!" << endl;
        for (auto& [k, v] : leaks) {
            ss << "{" << v.count << "} " << v.allocSize << " bytes at 0x" << k << endl;
        }
        ss << "Leak dump end." << endl;
        return ss.str();
    }

    // IMallocSpyを実装. Pre/Post, Alloc/Free 系の中で割当情報を保存しておく.
    virtual SIZE_T STDMETHODCALLTYPE PreAlloc(_In_ SIZE_T cbRequest) override
    {
        m_allocInfo.set(cbRequest);
        return cbRequest;
    }
    virtual void *STDMETHODCALLTYPE PostAlloc(_In_ void *pActual) override
    {
        if (pActual) {
            m_alloced[pActual] = m_allocInfo;
        }
        return pActual;
    }
    virtual void *STDMETHODCALLTYPE PreFree(_In_ void *pRequest, _In_ BOOL fSpyed) override
    {
        m_alloced.erase(pRequest);
        return pRequest;
    }
    virtual void STDMETHODCALLTYPE PostFree(_In_ BOOL fSpyed) override { }
    virtual SIZE_T STDMETHODCALLTYPE PreRealloc(_In_ void *pRequest, _In_ SIZE_T cbRequest, _Outptr_ void **ppNewRequest, _In_ BOOL fSpyed)
    {
        if (pRequest) {
            m_alloced.erase(pRequest);
        }
        m_allocInfo.set(cbRequest);
        return cbRequest;
    }
    virtual void *STDMETHODCALLTYPE PostRealloc(_In_ void *pActual, _In_ BOOL fSpyed)
    {
        if (pActual) {
            m_alloced[pActual] = m_allocInfo;
        }
        return pActual;
    }
    virtual void *STDMETHODCALLTYPE PreGetSize(_In_ void *pRequest, _In_ BOOL fSpyed)
    {
        return pRequest;
    }
    virtual SIZE_T STDMETHODCALLTYPE PostGetSize(_In_ SIZE_T cbActual, _In_ BOOL fSpyed) { return cbActual; }
    virtual void *STDMETHODCALLTYPE PreDidAlloc(_In_ void *pRequest, _In_ BOOL fSpyed) override { return pRequest; }
    virtual int STDMETHODCALLTYPE PostDidAlloc(_In_ void *pRequest, _In_ BOOL fSpyed, _In_ int fActual) override { return fActual; }
    virtual void STDMETHODCALLTYPE PreHeapMinimize(void) override { }
    virtual void STDMETHODCALLTYPE PostHeapMinimize(void) override { }
    // IUnknown を適当に実装.
    virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppvObject)
    {
        if (riid == IID_IMallocSpy) {
            *ppvObject = this;
            return S_OK;
        }
        return E_NOTIMPL;
    }
    virtual ULONG STDMETHODCALLTYPE AddRef(void)
    {
        return 1;
    }
    virtual ULONG STDMETHODCALLTYPE Release(void)
    {
        return 0;
    }
};

int main() {
    MallocSpy spy;
    // プロセスに登録
    ::CoRegisterMallocSpy(&spy);

    // p2だけリークさせる
    auto p = CoTaskMemAlloc(100);
    auto p2 = CoTaskMemAlloc(200);
    CoTaskMemFree(p);

    // リーク情報をダンプ
    cout << spy.DumpLeak() << endl;
}

出力はこうなる。

CoTaskMem Leak!
{2} 200 bytes at 0x00780C98
Leak dump end.

あとはGoogle Test のEventListenerなどの仕組みを使って単体テストの前後でリークチェックするなどしましょう。