mem_group.c

/*
 * Simple power-of-2 Allocate only mem-group for use with scatter-gather lists with read-write protection 
 * on the objects allocated from the memory group.
 */

#include<stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/mman.h>
#include "mem_group.h"

void memGroupInit(MemGroupT* obj,long long baseSize)
{
    long long i;
    char template[] = "/tmp/rowr.XXXXXX";
    int fd = mkstemp(template);
    assert(fd >= 0);
    strncpy(obj->tmpfile, template, sizeof(obj->tmpfile)-1);
    obj->tmpfd = fd;
    for (i =0;i<MEM_GRP_ARRAY_SIZE;i++)
    {
        obj->rwbuffers[i] = NULL;
        obj->robuffers[i] = NULL;
        obj->bufPos[i] = 0;
    }
    obj->fundamentalSize = baseSize;
    obj->curBuffer = -1;    
}

static char *memAlloc(MemGroupT *obj, long long size, int prot)
{
    char *map;
    int map_prot = PROT_READ;
    prot &= O_ACCMODE;
    if(prot & ( 1 << O_RDONLY)) 
        map_prot |= PROT_READ;
    if(prot & ( 1 << O_WRONLY)) 
        map_prot |= PROT_WRITE;
    if(obj->tmpfd < 0)
        map = mmap(0, size, map_prot, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
    else
        map = mmap(0, size, map_prot, MAP_SHARED, obj->tmpfd, 0);
    if(map == MAP_FAILED)
        return NULL;
    if(obj->tmpFileSize < size)
    {
        obj->tmpFileSize = size;
        ftruncate(obj->tmpfd, size);
    }
    return map;
}

static __inline__ void memFree(char *map, long long size)
{
    munmap(map, size);
}

void* memGroupAlloc(MemGroupT* obj,int amt)
{
    void* ret;
    int i;

    while(obj->curBuffer<MEM_GRP_ARRAY_SIZE)
    {
        
        long long bSize = obj->fundamentalSize<<obj->curBuffer;
        /* Look through the allocated buffers to see if there is room for what
           the caller wants at the end of one of them */
        for (i = obj->curBuffer;i>=0 && bSize >= amt; i--,bSize>>=1)
        {
            if (bSize-obj->bufPos[i]>=amt)
            {
                ret = obj->rwbuffers[i]+obj->bufPos[i];
                obj->bufPos[i] += amt;
                return ret;
            }
        }
        /*
         * Check for the max limit.
         */
        if(obj->curBuffer + 1 == MEM_GRP_ARRAY_SIZE)
            return NULL;
        /* Nothing found so we have to allocate a new buffer */
        obj->curBuffer+=1;
        obj->bufPos[obj->curBuffer] = 0;
        /* Each time thru I allocate 2x the memory */
        obj->rwbuffers[obj->curBuffer] = memAlloc(obj, obj->fundamentalSize<<obj->curBuffer, O_RDWR);
        if (!obj->rwbuffers[obj->curBuffer]) return NULL; /* no memory to be had */   
        obj->robuffers[obj->curBuffer] = memAlloc(obj, obj->fundamentalSize<<obj->curBuffer, O_RDONLY);
        if(!obj->robuffers[obj->curBuffer]) return NULL;
    }
    /* Memory Group is completely full */
    return NULL; /* no memory to be had */    
}

void memGroupFreeAll(MemGroupT* obj)
{
    long long i;
    
    for (i = obj->curBuffer;i>=0;i--)
    {
        memFree(obj->rwbuffers[i], obj->fundamentalSize << i);
        memFree(obj->robuffers[i], obj->fundamentalSize << i);
        obj->rwbuffers[i] = NULL;
        obj->robuffers[i] = NULL;
        obj->bufPos[i] = 0;            
    }
    obj->curBuffer = -1;
}

static void *memGroupFind(MemGroupT *grp, char *obj, char **srcBuffer, int *index)
{
    register int i;
    for(i = grp->curBuffer; i >= 0; --i)
    {
        long long size;
        if(obj < srcBuffer[i]) continue;
        size = grp->fundamentalSize << i;
        if(obj > srcBuffer[i] + size) continue;
        if(index)
            *index = i;
        return srcBuffer[i];
    }
    return NULL;
}

char *memGroupProtect(MemGroupT *grp, char *obj)
{
    long long offset = 0;
    int index = 0;
    char *base = memGroupFind(grp, obj, grp->rwbuffers, &index);
    if(!base) return NULL;
    offset = obj - base; 
    /*
     * return the offset to the read only buffer.
     */
    return grp->robuffers[index] + offset;
}

char *memGroupUnprotect(MemGroupT *grp, char *obj)
{
    char *base = NULL;
    int index = 0;
    long long offset;
    base = memGroupFind(grp, obj, grp->robuffers, &index);
    if(!base) return NULL;
    offset = obj - base;
    return grp->rwbuffers[index] + offset;
}

void memGroupDelete(MemGroupT* obj)
{
    memGroupFreeAll(obj);
    if(obj->tmpfd >= 0)
        close(obj->tmpfd);
    if(obj->tmpfile[0])
        unlink(obj->tmpfile);
    obj->fundamentalSize = 0;
}

int main(int argc, char **argv)
{
    MemGroupT memGroup = {{0}};
    struct objStorage
    {
        char **objs;
        long long *sizes;
        int count;
    } objStorage = { .objs = NULL, .count = 10 , .sizes = NULL};
    register int i;
    /*
     * setup the sigsegv handler.
     */
    memGroupInit(&memGroup, 16);
    if(argc > 1)
    {
        char *p = NULL;
        objStorage.count = (int) strtol(argv[1], &p, 10);
        if(*p) 
        {
            fprintf(stderr, "Invalid count specified [%s]\n", argv[1]);
            return -1;
        }
        if(!objStorage.count) objStorage.count = 10;
    }
    objStorage.objs = calloc(objStorage.count, sizeof(*objStorage.objs));
    assert(objStorage.objs);
    objStorage.sizes = calloc(objStorage.count, sizeof(*objStorage.sizes));
    for(i = 0; i < objStorage.count; ++i)
    {
        long long size = random() & 1023;
        if(!size) size = 100;
        objStorage.sizes[i]= size;
        objStorage.objs[i] = memGroupAlloc(&memGroup, size);
        assert(objStorage.objs[i]);
        memset(objStorage.objs[i], 0xaa, size);
    }
    fprintf(stderr, "Written [%d] stuffs to the mem group\n", objStorage.count);
    /*
     * Now unprotect all of them.
     */
    for(i = 0; i < objStorage.count; ++i)
    {
        char *obj = objStorage.objs[i];
        char *restoreobj;
        objStorage.objs[i] = memGroupProtect(&memGroup, obj);
        assert(objStorage.objs[i] != obj);
        if(memcmp(objStorage.objs[i], obj, objStorage.sizes[i]))
        {
            fprintf(stderr, "Obj compare failed for index [%d]\n", i);
            assert(0);
        }
        /*
         * trigger a SIGSEGV and wait for the trigger
         */
#if 0
        *objStorage.objs[i] = 0xaa;
        while(!got_sigsegv) sleep(1);
        got_sigsegv = 0;
#endif
        restoreobj = objStorage.objs[i];
        objStorage.objs[i] = memGroupUnprotect(&memGroup, restoreobj);
        assert(objStorage.objs[i] != restoreobj);
        if(memcmp(objStorage.objs[i], restoreobj, objStorage.sizes[i]))
        {
            fprintf(stderr, "Obj2 compare failed for index [%d]\n", i);
            assert(0);
        }
        /*
         * Do a write to verify that the process doesn't core dump on a sigsegv.
         */
        memset(objStorage.objs[i], 0xaa, objStorage.sizes[i]);
    }
    memGroupDelete(&memGroup);
    free(objStorage.objs);
    free(objStorage.sizes);
    return 0;
}


/*
 * Local variables:
 * c-file-style: "linux"
 * compile-command: "gcc -Wall -g -o mem_group mem_group.c"
 * tab-width: 4
 * End:
*/

mem_group.h

#ifndef _MEM_GROUP_H_
#define _MEM_GROUP_H_

#include <stdio.h>

#ifdef __cplusplus
extern "C" {
#endif

#define MEM_GRP_ARRAY_SIZE 32

typedef struct MemGroup
{
    char tmpfile[0xff+1];
    int tmpfd;
    int tmpFileSize;
    /** An array of pointers to memory buffers that memory is served out of.  Each buffer is 2x the size of the prior */
    char*   rwbuffers[MEM_GRP_ARRAY_SIZE];
    char*   robuffers[MEM_GRP_ARRAY_SIZE];
    /** The beginning of the unallocated area of each buffer */
    long long bufPos[MEM_GRP_ARRAY_SIZE];
    /** The size of the first buffer.  This should be set to some multiple of the size of your typical allocation. */
    long long fundamentalSize;
    /** The last allocated buffer */
    int     curBuffer; 
} MemGroupT;

extern void memGroupInit(MemGroupT* obj,long long baseSize);
extern void* memGroupAlloc(MemGroupT* obj,int amt);
extern void memGroupFreeAll(MemGroupT* obj);
extern void memGroupDelete(MemGroupT* obj);
extern char *memGroupProtect(MemGroupT *grp, char *obj);
extern char *memGroupUnprotect(MemGroupT *grp, char *obj);   

#ifdef __cplusplus
}
#endif

#endif