Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
/*
* 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:
*/
#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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment