From d49aee3fa631eac9d5483c600e9553aa15442fde Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Tue, 6 Sep 2016 10:35:34 -0700 Subject: [PATCH] Add new flags for THMapAllocator * TH_ALLOCATOR_MAPPED_FROMFD uses an existing file descriptor for mapping (and steals it) * TH_ALLOCATOR_MAPPED_KEEPFD doesn't close the file descriptor until the mapping is freed * TH_ALLOCATOR_MAPPED_UNLINK unlinks the file immediately after mapping it to memory Also, now it's using fstat to check the file size (instead of lseek, which alters the fd state). --- lib/TH/THAllocator.c | 137 ++++++++++++++++++++++++++++++++----------- lib/TH/THAllocator.h | 6 ++ 2 files changed, 108 insertions(+), 35 deletions(-) diff --git a/lib/TH/THAllocator.c b/lib/TH/THAllocator.c index fea6d4bf..abfd3a6c 100644 --- a/lib/TH/THAllocator.c +++ b/lib/TH/THAllocator.c @@ -39,6 +39,7 @@ struct THMapAllocatorContext_ { char *filename; /* file name */ int flags; long size; /* mapped size */ + int fd; }; #define TH_ALLOC_ALIGNMENT 64 @@ -47,21 +48,35 @@ typedef struct { int refcount; } THMapInfo; +char * unknown_filename = "filename not specified"; + THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags) { THMapAllocatorContext *ctx = THAlloc(sizeof(THMapAllocatorContext)); - if (!(flags & TH_ALLOCATOR_MAPPED_SHARED) && !(flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)) flags &= ~TH_ALLOCATOR_MAPPED_NOCREATE; if ((flags ^ TH_ALLOCATOR_MAPPED_EXCLUSIVE) == 0) THError("TH_ALLOCATOR_MAPPED_EXCLUSIVE flag requires opening the file " "in shared mode"); - ctx->filename = THAlloc(strlen(filename)+1); - strcpy(ctx->filename, filename); + if (filename) { + ctx->filename = THAlloc(strlen(filename)+1); + strcpy(ctx->filename, filename); + } else { + ctx->filename = unknown_filename; + } ctx->flags = flags; ctx->size = 0; + ctx->fd = -1; + + return ctx; +} + +THMapAllocatorContext *THMapAllocatorContext_newWithFd(const char *filename, int fd, int flags) +{ + THMapAllocatorContext *ctx = THMapAllocatorContext_new(filename, flags); + ctx->fd = fd; return ctx; } @@ -71,6 +86,11 @@ char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx) return ctx->filename; } +int THMapAllocatorContext_fd(THMapAllocatorContext *ctx) +{ + return ctx->fd; +} + long THMapAllocatorContext_size(THMapAllocatorContext *ctx) { return ctx->size; @@ -78,7 +98,8 @@ long THMapAllocatorContext_size(THMapAllocatorContext *ctx) void THMapAllocatorContext_free(THMapAllocatorContext *ctx) { - THFree(ctx->filename); + if (ctx->filename != unknown_filename) + THFree(ctx->filename); THFree(ctx); } @@ -98,6 +119,10 @@ static void *_map_alloc(void* ctx_, long size) THError("exclusive file mapping is not supported on Windows"); if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE) THError("file mapping without creation is not supported on Windows"); + if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) + THError("TH_ALLOCATOR_MAPPED_KEEPFD not supported on Windows"); + if (ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD) + THError("TH_ALLOCATOR_MAPPED_FROMFD not supported on Windows"); /* open file */ /* FILE_FLAG_RANDOM_ACCESS ? */ @@ -193,9 +218,9 @@ static void *_map_alloc(void* ctx_, long size) /* open file */ int fd; int flags; - long fdsz; + struct stat file_stat; - if (ctx->flags) + if (ctx->flags & (TH_ALLOCATOR_MAPPED_SHARED | TH_ALLOCATOR_MAPPED_SHAREDMEM)) flags = O_RDWR | O_CREAT; else flags = O_RDONLY; @@ -205,35 +230,40 @@ static void *_map_alloc(void* ctx_, long size) if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE) flags &= ~O_CREAT; - if(ctx->flags & TH_ALLOCATOR_MAPPED_SHARED) - { - if((fd = open(ctx->filename, flags, (mode_t)0600)) == -1) - THError("unable to open file <%s> in read-write mode", ctx->filename); - } - else if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM) - { + if (!(ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD)) { + if(ctx->flags & TH_ALLOCATOR_MAPPED_SHARED) + { + if((fd = open(ctx->filename, flags, (mode_t)0600)) == -1) + THError("unable to open file <%s> in read-write mode", ctx->filename); + } + else if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM) + { #ifdef HAVE_SHM_OPEN - if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1) - THError("unable to open shared memory object <%s> in read-write mode", ctx->filename); + if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1) + THError("unable to open shared memory object <%s> in read-write mode", ctx->filename); #else - THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform"); + THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform"); #endif - } - else - { - if((fd = open(ctx->filename, O_RDONLY)) == -1) - THError("unable to open file <%s> in read-only mode", ctx->filename); + } + else + { + if((fd = open(ctx->filename, O_RDONLY)) == -1) + THError("unable to open file <%s> in read-only mode", ctx->filename); + } + } else { + fd = ctx->fd; } - if((fdsz = lseek(fd, 0, SEEK_END)) == -1) + if(fstat(fd, &file_stat) == -1) { - close(fd); - THError("unable to seek at end of file <%s>", ctx->filename); + if (!(ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD)) + close(fd); + THError("unable to stat the file <%s>", ctx->filename); } if(size > 0) { - if(size > fdsz) + if(size > file_stat.st_size) { if(ctx->flags) { @@ -243,7 +273,7 @@ static void *_map_alloc(void* ctx_, long size) if(ftruncate(fd, size) == -1) THError("unable to resize shared memory file <%s> to the right size", ctx->filename); } - if((fdsz = lseek(fd, size-1, SEEK_SET)) == -1) + if(fstat(fd, &file_stat) == -1 || file_stat.st_size < size) { close(fd); THError("unable to stretch file <%s> to the right size", ctx->filename); @@ -262,18 +292,40 @@ static void *_map_alloc(void* ctx_, long size) } } else - size = fdsz; + size = file_stat.st_size; ctx->size = size; /* if we are here, it must be the right size */ /* map it */ - if(ctx->flags) + if (ctx->flags & (TH_ALLOCATOR_MAPPED_SHARED | TH_ALLOCATOR_MAPPED_SHAREDMEM)) data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0); else data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0); - if(close(fd) == -1) - THError("Error closing file <%s>", ctx->filename); + if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) { + ctx->fd = fd; + } else { + if(close(fd) == -1) + THError("Error closing file <%s>", ctx->filename); + ctx->fd = -1; + } + + if (ctx->flags & TH_ALLOCATOR_MAPPED_UNLINK) { + if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM) + { +#ifdef HAVE_SHM_UNLINK + if (shm_unlink(ctx->filename) == -1) + THError("could not unlink the shared memory file %s", ctx->filename); +#else + THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename); +#endif + } + else + { + if (unlink(ctx->filename) == -1) + THError("could not unlink file %s", ctx->filename); + } + } if(data == MAP_FAILED) { @@ -302,16 +354,25 @@ static void THMapAllocator_free(void* ctx_, void* data) { if(!UnmapViewOfFile((LPINT)data)) THError("could not unmap the shared memory file"); #else /* _WIN32 */ + if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) { + if (close(ctx->fd) == -1) + THError("could not close file descriptor %d", ctx->fd); + } + if (munmap(data, ctx->size)) THError("could not unmap the shared memory file"); - if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM) + + if (!(ctx->flags & (TH_ALLOCATOR_MAPPED_FROMFD | TH_ALLOCATOR_MAPPED_UNLINK))) { + if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM) + { #ifdef HAVE_SHM_UNLINK - if (shm_unlink(ctx->filename) == -1) - THError("could not unlink the shared memory file %s", ctx->filename); + if (shm_unlink(ctx->filename) == -1) + THError("could not unlink the shared memory file %s", ctx->filename); #else - THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename); + THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename); #endif + } } #endif /* _WIN32 */ @@ -350,8 +411,14 @@ static void THMapAllocator_free(void* ctx, void* data) { static void * THRefcountedMapAllocator_alloc(void *_ctx, long size) { THMapAllocatorContext *ctx = _ctx; + if (ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD) + THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_FROMFD flag"); + if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) + THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_KEEPFD flag"); + if (ctx->flags & TH_ALLOCATOR_MAPPED_UNLINK) + THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_UNLINK flag"); if (!(ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)) - THError("THRefcountedMapAllcator requires SHAREDMEM flag"); + THError("THRefcountedMapAllocator requires TH_ALLOCATOR_MAPPED_SHAREDMEM flag"); size = size + TH_ALLOC_ALIGNMENT; void *ptr = _map_alloc(ctx, size); diff --git a/lib/TH/THAllocator.h b/lib/TH/THAllocator.h index 567c59f8..049d08dd 100644 --- a/lib/TH/THAllocator.h +++ b/lib/TH/THAllocator.h @@ -7,6 +7,9 @@ #define TH_ALLOCATOR_MAPPED_SHAREDMEM 2 #define TH_ALLOCATOR_MAPPED_EXCLUSIVE 4 #define TH_ALLOCATOR_MAPPED_NOCREATE 8 +#define TH_ALLOCATOR_MAPPED_KEEPFD 16 +#define TH_ALLOCATOR_MAPPED_FROMFD 32 +#define TH_ALLOCATOR_MAPPED_UNLINK 64 /* Custom allocator */ @@ -25,7 +28,10 @@ extern THAllocator THDefaultAllocator; */ typedef struct THMapAllocatorContext_ THMapAllocatorContext; TH_API THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags); +TH_API THMapAllocatorContext *THMapAllocatorContext_newWithFd(const char *filename, + int fd, int flags); TH_API char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx); +TH_API int THMapAllocatorContext_fd(THMapAllocatorContext *ctx); TH_API long THMapAllocatorContext_size(THMapAllocatorContext *ctx); TH_API void THMapAllocatorContext_free(THMapAllocatorContext *ctx);