malloc.cpp

来自「firtext搜索引擎源码」· C++ 代码 · 共 618 行 · 第 1/2 页

#include "com/comerror.h"
#include "com/Com.h"
#include <assert.h>
#include <pthread.h>
#include <stdlib.h>
#include "com/linux/internal.h"

namespace firtex
{
	namespace com 
	{
		namespace internal 
		{
			class Malloc : public IMalloc 
			{
			public:
				Malloc();
				~Malloc();

				FX_STDMETHOD(QueryInterface)( firtex::com::FX_IID const& riid, void** ptr );
				FX_STDMETHOD_(uint32_t,AddRef)();
				FX_STDMETHOD_(uint32_t,Release)();
				FX_STDMETHOD_(void*,Alloc)(uint32_t);
				FX_STDMETHOD_(void*,Realloc)(void*,uint32_t);
				FX_STDMETHOD_(void,Free)(void*);
				FX_STDMETHOD_(uint32_t,GetSize)(void*);
				FX_STDMETHOD_(int16_t,DidAlloc)(void*);
				FX_STDMETHOD_(void,HeapMinimize)();

			public:
				IMallocSpy* pSpy;          /* the spy when active */
				uint32_t SpyedAllocationsLeft; /* number of spyed allocations left */
				bool SpyReleasePending;     /* FX_CoRevokeMallocSpy called with spyed allocations left*/
				void** SpyedBlocks;       /* root of the table */
				int SpyedBlockTableLength;  /* size of the table*/
			};
		}
	}
}

static firtex::com::internal::Malloc malloc32;
static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
	
/* resize the old table */
static firtex::com:: FX_HRESULT SetSpyedBlockTableLength ( int NewLength )
{
	void** NewSpyedBlocks = 0;

	if (! malloc32.SpyedBlocks ) {
		NewSpyedBlocks = (void**)malloc(NewLength*sizeof(void*));
		memset( NewSpyedBlocks, 0, NewLength*sizeof(void*) );
	}
	else {
		NewSpyedBlocks = (void**)realloc( malloc32.SpyedBlocks, NewLength*sizeof(void*));
		if ( NewLength > malloc32.SpyedBlockTableLength )
			memset( NewSpyedBlocks + malloc32.SpyedBlockTableLength, 0, (NewLength - malloc32.SpyedBlockTableLength)*sizeof(void*) );
	}

	if (NewSpyedBlocks) {
		malloc32.SpyedBlocks = NewSpyedBlocks;
		malloc32.SpyedBlockTableLength = NewLength;
		return FX_S_OK;
	}

	return FX_E_OUTOFMEMORY;
}

/* add a location to the table */
static firtex::com::FX_HRESULT AddMemoryLocation(void* pMem)
{
	/* allocate the table if not already allocated */
	if ( !malloc32.SpyedBlockTableLength ) {
		firtex::com::FX_HRESULT hr = SetSpyedBlockTableLength( 0x1000 );
		if ( FX_FAILED(hr) ) return hr;
	}

	/* find a free location */
	void** Current = malloc32.SpyedBlocks;
	while (*Current) {
		++Current;

	    if ( Current >= malloc32.SpyedBlocks + malloc32.SpyedBlockTableLength ) {
	    	/* no more space in table, grow it */
			firtex::com::FX_HRESULT hr = SetSpyedBlockTableLength( malloc32.SpyedBlockTableLength + 0x1000 );
			if ( FX_FAILED(hr) ) return hr;
	    }
	};

	/* put the location in our table */
	*Current = pMem;
	++malloc32.SpyedAllocationsLeft;
	return FX_S_OK;
}

static firtex::com::FX_HRESULT RemoveMemoryLocation(void* pMem)
{
	/* allocate the table if not already allocated */
	if ( !malloc32.SpyedBlockTableLength ) {
		firtex::com::FX_HRESULT hr = SetSpyedBlockTableLength( 0x1000 );
		if ( FX_FAILED(hr) ) return hr;
	}

	void** Current = malloc32.SpyedBlocks;

	/* find the location */
	while (*Current != pMem) {
		++Current;
	    if (Current >= malloc32.SpyedBlocks + malloc32.SpyedBlockTableLength)  
			return FX_E_FAIL;      /* not found  */
	}

	/* location found */
    --malloc32.SpyedAllocationsLeft;
	/*TRACE("%lu\n",Malloc32.SpyedAllocationsLeft);*/
	*Current = NULL;
	return FX_S_OK;
}

firtex::com::internal::Malloc::Malloc()
{
	pSpy = 0;          
	SpyedAllocationsLeft = 0; 
	SpyReleasePending = false;
	
	SpyedBlocks = 0;
	SpyedBlockTableLength = 0;
}

firtex::com::internal::Malloc::~Malloc()
{
	free( SpyedBlocks );
}


/******************************************************************************
 *	IMalloc32_QueryInterface	[VTABLE]
 */
firtex::com::FX_HRESULT firtex::com::internal::Malloc::QueryInterface(FX_REFIID refiid,void** obj) 
{
	if ( !obj ) return FX_E_INVALIDARG;

	if ( FX_IsEqualIID(IUnknown::iid,refiid) ) {
		*obj = static_cast<IUnknown*>(this);
		return FX_S_OK;
	}
	else if ( FX_IsEqualIID(IMalloc::iid,refiid) ) {
		*obj = static_cast<IMalloc*>(this);
		return FX_S_OK;
	}
	return FX_E_NOINTERFACE;
}

/******************************************************************************
 *	IMalloc32_AddRefRelease		[VTABLE]
 */
uint32_t firtex::com::internal::Malloc::AddRef() {
	return 1;
}

uint32_t firtex::com::internal::Malloc::Release() {
	return 1;
}

/******************************************************************************
 *	IMalloc32_Alloc 		[VTABLE]
 */
void* firtex::com::internal::Malloc::Alloc(uint32_t cb) 
{
	if ( malloc32.pSpy ) {
	    uint32_t preAllocResult;

		::pthread_mutex_lock( &mutex );
	    
	    preAllocResult = malloc32.pSpy->PreAlloc(cb);
	    if ((cb != 0) && (preAllocResult == 0)) {
			/* PreAlloc can force Alloc to fail, but not if cb == 0 */
			::pthread_mutex_unlock( &mutex );
			return NULL;
	    }
	}
 	
	void* addr = malloc(cb);

	if ( malloc32.pSpy ) {
	    addr = malloc32.pSpy->PostAlloc(addr);
	    if (addr) AddMemoryLocation(addr);
		::pthread_mutex_unlock( &mutex );
	}

	return addr;
}

/******************************************************************************
 * IMalloc32_Realloc [VTABLE]
 */
void* firtex::com::internal::Malloc::Realloc(void* pv,uint32_t cb) 
{
	if ( malloc32.pSpy ) {
	    void* pRealMemory;
	    bool fSpyed;

		::pthread_mutex_lock( &mutex );

		FX_HRESULT hr = RemoveMemoryLocation(pv);
		fSpyed = FX_SUCCEEDED(hr);
		cb = malloc32.pSpy->PreRealloc( pv, cb, &pRealMemory, fSpyed);

	    /* check if can release the spy */
	    if( malloc32.SpyReleasePending && !malloc32.SpyedAllocationsLeft ) {
			malloc32.SpyReleasePending = false;
			malloc32.pSpy->Release();
			malloc32.pSpy = 0;
	    }

	    if (0==cb) {
	        /* PreRealloc can force Realloc to fail */
			::pthread_mutex_unlock( &mutex );
			return NULL;
	    }
	    pv = pRealMemory;
	}

	void* pNewMemory = 0;
    if (!pv) 
		pNewMemory = malloc(cb);
	else if (cb) 
		pNewMemory = realloc(pv,cb);
	else {
		free( pv );
	    pNewMemory = NULL;
	}

	if( malloc32.pSpy ) {
	    pNewMemory = malloc32.pSpy->PostRealloc(pNewMemory, -1);
	    if (pNewMemory) 
			AddMemoryLocation(pNewMemory);
		::pthread_mutex_unlock( &mutex );
	}

	return pNewMemory;
}

/******************************************************************************
 * IMalloc32_Free [VTABLE]
 */
void firtex::com::internal::Malloc::Free(void* pv) 
{
	bool fSpyed = false;

	if ( malloc32.pSpy ) {
		::pthread_mutex_lock( &mutex );
  		FX_HRESULT hr = RemoveMemoryLocation(pv);
		fSpyed = FX_SUCCEEDED(hr);

	    pv = malloc32.pSpy->PreFree(pv, fSpyed);
	}

	free( pv );

	if( malloc32.pSpy ) {
		malloc32.pSpy->PostFree( fSpyed );

	    /* check if can release the spy */
	    if(malloc32.SpyReleasePending && !malloc32.SpyedAllocationsLeft) {
			malloc32.SpyReleasePending = false;
			malloc32.pSpy->Release();
			malloc32.pSpy = NULL;
	    }

		::pthread_mutex_unlock( &mutex );
	}
}

/******************************************************************************
 * IMalloc32_GetSize [VTABLE]
 *
 * NOTES
 *  FIXME returns:
 *      win95:  size allocated (4 byte boundarys)
 *      win2k:  size originally requested !!! (allocated on 8 byte boundarys)
 */
uint32_t firtex::com::internal::Malloc::GetSize(void* pv) 
{
	bool fSpyed = false;

	if(malloc32.pSpy) {
		::pthread_mutex_lock( &mutex );
	    pv = malloc32.pSpy->PreGetSize(pv, fSpyed);
	}

	// oh, oh, what do we do here
	//cb = HeapSize(GetProcessHeap(),0,pv);
	//?uint32_t cb = msize( pv );
	uint32_t cb = 0;

	if( malloc32.pSpy ) {
	    cb = malloc32.pSpy->PostGetSize(cb, fSpyed);
		::pthread_mutex_unlock( &mutex );
	}

	return cb;
}

/******************************************************************************
 * IMalloc32_DidAlloc [VTABLE]
 */
int16_t firtex::com::internal::Malloc::DidAlloc(void* pv) 
{
	bool fSpyed = false;