Skip to content

Commit 7bbe179

Browse files
soumithsoumith
authored
Merge pull request #757 from apaszke/master
Add new flags for THMapAllocator
2 parents 649de92 + d49aee3 commit 7bbe179

File tree

2 files changed

+108
-35
lines changed

2 files changed

+108
-35
lines changed

‎lib/TH/THAllocator.c

Lines changed: 102 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ struct THMapAllocatorContext_ {
3939
char *filename; /* file name */
4040
int flags;
4141
long size; /* mapped size */
42+
int fd;
4243
};
4344

4445
#define TH_ALLOC_ALIGNMENT 64
@@ -47,21 +48,35 @@ typedef struct {
4748
int refcount;
4849
} THMapInfo;
4950

51+
char * unknown_filename = "filename not specified";
52+
5053
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags)
5154
{
5255
THMapAllocatorContext *ctx = THAlloc(sizeof(THMapAllocatorContext));
5356

54-
5557
if (!(flags & TH_ALLOCATOR_MAPPED_SHARED) && !(flags & TH_ALLOCATOR_MAPPED_SHAREDMEM))
5658
flags &= ~TH_ALLOCATOR_MAPPED_NOCREATE;
5759
if ((flags ^ TH_ALLOCATOR_MAPPED_EXCLUSIVE) == 0)
5860
THError("TH_ALLOCATOR_MAPPED_EXCLUSIVE flag requires opening the file "
5961
"in shared mode");
6062

61-
ctx->filename = THAlloc(strlen(filename)+1);
62-
strcpy(ctx->filename, filename);
63+
if (filename) {
64+
ctx->filename = THAlloc(strlen(filename)+1);
65+
strcpy(ctx->filename, filename);
66+
} else {
67+
ctx->filename = unknown_filename;
68+
}
6369
ctx->flags = flags;
6470
ctx->size = 0;
71+
ctx->fd = -1;
72+
73+
return ctx;
74+
}
75+
76+
THMapAllocatorContext *THMapAllocatorContext_newWithFd(const char *filename, int fd, int flags)
77+
{
78+
THMapAllocatorContext *ctx = THMapAllocatorContext_new(filename, flags);
79+
ctx->fd = fd;
6580

6681
return ctx;
6782
}
@@ -71,14 +86,20 @@ char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx)
7186
return ctx->filename;
7287
}
7388

89+
int THMapAllocatorContext_fd(THMapAllocatorContext *ctx)
90+
{
91+
return ctx->fd;
92+
}
93+
7494
long THMapAllocatorContext_size(THMapAllocatorContext *ctx)
7595
{
7696
return ctx->size;
7797
}
7898

7999
void THMapAllocatorContext_free(THMapAllocatorContext *ctx)
80100
{
81-
THFree(ctx->filename);
101+
if (ctx->filename != unknown_filename)
102+
THFree(ctx->filename);
82103
THFree(ctx);
83104
}
84105

@@ -98,6 +119,10 @@ static void *_map_alloc(void* ctx_, long size)
98119
THError("exclusive file mapping is not supported on Windows");
99120
if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE)
100121
THError("file mapping without creation is not supported on Windows");
122+
if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD)
123+
THError("TH_ALLOCATOR_MAPPED_KEEPFD not supported on Windows");
124+
if (ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD)
125+
THError("TH_ALLOCATOR_MAPPED_FROMFD not supported on Windows");
101126

102127
/* open file */
103128
/* FILE_FLAG_RANDOM_ACCESS ? */
@@ -193,9 +218,9 @@ static void *_map_alloc(void* ctx_, long size)
193218
/* open file */
194219
int fd;
195220
int flags;
196-
long fdsz;
221+
struct stat file_stat;
197222

198-
if (ctx->flags)
223+
if (ctx->flags & (TH_ALLOCATOR_MAPPED_SHARED | TH_ALLOCATOR_MAPPED_SHAREDMEM))
199224
flags = O_RDWR | O_CREAT;
200225
else
201226
flags = O_RDONLY;
@@ -205,35 +230,40 @@ static void *_map_alloc(void* ctx_, long size)
205230
if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE)
206231
flags &= ~O_CREAT;
207232

208-
if(ctx->flags & TH_ALLOCATOR_MAPPED_SHARED)
209-
{
210-
if((fd = open(ctx->filename, flags, (mode_t)0600)) == -1)
211-
THError("unable to open file <%s> in read-write mode", ctx->filename);
212-
}
213-
else if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
214-
{
233+
if (!(ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD)) {
234+
if(ctx->flags & TH_ALLOCATOR_MAPPED_SHARED)
235+
{
236+
if((fd = open(ctx->filename, flags, (mode_t)0600)) == -1)
237+
THError("unable to open file <%s> in read-write mode", ctx->filename);
238+
}
239+
else if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
240+
{
215241
#ifdef HAVE_SHM_OPEN
216-
if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1)
217-
THError("unable to open shared memory object <%s> in read-write mode", ctx->filename);
242+
if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1)
243+
THError("unable to open shared memory object <%s> in read-write mode", ctx->filename);
218244
#else
219-
THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform");
245+
THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform");
220246
#endif
221-
}
222-
else
223-
{
224-
if((fd = open(ctx->filename, O_RDONLY)) == -1)
225-
THError("unable to open file <%s> in read-only mode", ctx->filename);
247+
}
248+
else
249+
{
250+
if((fd = open(ctx->filename, O_RDONLY)) == -1)
251+
THError("unable to open file <%s> in read-only mode", ctx->filename);
252+
}
253+
} else {
254+
fd = ctx->fd;
226255
}
227256

228-
if((fdsz = lseek(fd, 0, SEEK_END)) == -1)
257+
if(fstat(fd, &file_stat) == -1)
229258
{
230-
close(fd);
231-
THError("unable to seek at end of file <%s>", ctx->filename);
259+
if (!(ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD))
260+
close(fd);
261+
THError("unable to stat the file <%s>", ctx->filename);
232262
}
233263

234264
if(size > 0)
235265
{
236-
if(size > fdsz)
266+
if(size > file_stat.st_size)
237267
{
238268
if(ctx->flags)
239269
{
@@ -243,7 +273,7 @@ static void *_map_alloc(void* ctx_, long size)
243273
if(ftruncate(fd, size) == -1)
244274
THError("unable to resize shared memory file <%s> to the right size", ctx->filename);
245275
}
246-
if((fdsz = lseek(fd, size-1, SEEK_SET)) == -1)
276+
if(fstat(fd, &file_stat) == -1 || file_stat.st_size < size)
247277
{
248278
close(fd);
249279
THError("unable to stretch file <%s> to the right size", ctx->filename);
@@ -262,18 +292,40 @@ static void *_map_alloc(void* ctx_, long size)
262292
}
263293
}
264294
else
265-
size = fdsz;
295+
size = file_stat.st_size;
266296

267297
ctx->size = size; /* if we are here, it must be the right size */
268298

269299
/* map it */
270-
if(ctx->flags)
300+
if (ctx->flags & (TH_ALLOCATOR_MAPPED_SHARED | TH_ALLOCATOR_MAPPED_SHAREDMEM))
271301
data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
272302
else
273303
data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0);
274304

275-
if(close(fd) == -1)
276-
THError("Error closing file <%s>", ctx->filename);
305+
if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) {
306+
ctx->fd = fd;
307+
} else {
308+
if(close(fd) == -1)
309+
THError("Error closing file <%s>", ctx->filename);
310+
ctx->fd = -1;
311+
}
312+
313+
if (ctx->flags & TH_ALLOCATOR_MAPPED_UNLINK) {
314+
if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
315+
{
316+
#ifdef HAVE_SHM_UNLINK
317+
if (shm_unlink(ctx->filename) == -1)
318+
THError("could not unlink the shared memory file %s", ctx->filename);
319+
#else
320+
THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename);
321+
#endif
322+
}
323+
else
324+
{
325+
if (unlink(ctx->filename) == -1)
326+
THError("could not unlink file %s", ctx->filename);
327+
}
328+
}
277329

278330
if(data == MAP_FAILED)
279331
{
@@ -302,16 +354,25 @@ static void THMapAllocator_free(void* ctx_, void* data) {
302354
if(!UnmapViewOfFile((LPINT)data))
303355
THError("could not unmap the shared memory file");
304356
#else /* _WIN32 */
357+
if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD) {
358+
if (close(ctx->fd) == -1)
359+
THError("could not close file descriptor %d", ctx->fd);
360+
}
361+
305362
if (munmap(data, ctx->size))
306363
THError("could not unmap the shared memory file");
307-
if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
364+
365+
if (!(ctx->flags & (TH_ALLOCATOR_MAPPED_FROMFD | TH_ALLOCATOR_MAPPED_UNLINK)))
308366
{
367+
if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
368+
{
309369
#ifdef HAVE_SHM_UNLINK
310-
if (shm_unlink(ctx->filename) == -1)
311-
THError("could not unlink the shared memory file %s", ctx->filename);
370+
if (shm_unlink(ctx->filename) == -1)
371+
THError("could not unlink the shared memory file %s", ctx->filename);
312372
#else
313-
THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename);
373+
THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename);
314374
#endif
375+
}
315376
}
316377
#endif /* _WIN32 */
317378

@@ -350,8 +411,14 @@ static void THMapAllocator_free(void* ctx, void* data) {
350411
static void * THRefcountedMapAllocator_alloc(void *_ctx, long size) {
351412
THMapAllocatorContext *ctx = _ctx;
352413

414+
if (ctx->flags & TH_ALLOCATOR_MAPPED_FROMFD)
415+
THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_FROMFD flag");
416+
if (ctx->flags & TH_ALLOCATOR_MAPPED_KEEPFD)
417+
THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_KEEPFD flag");
418+
if (ctx->flags & TH_ALLOCATOR_MAPPED_UNLINK)
419+
THError("THRefcountedMapAllocator doesn't support TH_ALLOCATOR_MAPPED_UNLINK flag");
353420
if (!(ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM))
354-
THError("THRefcountedMapAllcator requires SHAREDMEM flag");
421+
THError("THRefcountedMapAllocator requires TH_ALLOCATOR_MAPPED_SHAREDMEM flag");
355422

356423
size = size + TH_ALLOC_ALIGNMENT;
357424
void *ptr = _map_alloc(ctx, size);

‎lib/TH/THAllocator.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77
#define TH_ALLOCATOR_MAPPED_SHAREDMEM 2
88
#define TH_ALLOCATOR_MAPPED_EXCLUSIVE 4
99
#define TH_ALLOCATOR_MAPPED_NOCREATE 8
10+
#define TH_ALLOCATOR_MAPPED_KEEPFD 16
11+
#define TH_ALLOCATOR_MAPPED_FROMFD 32
12+
#define TH_ALLOCATOR_MAPPED_UNLINK 64
1013

1114
/* Custom allocator
1215
*/
@@ -25,7 +28,10 @@ extern THAllocator THDefaultAllocator;
2528
*/
2629
typedef struct THMapAllocatorContext_ THMapAllocatorContext;
2730
TH_API THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags);
31+
TH_API THMapAllocatorContext *THMapAllocatorContext_newWithFd(const char *filename,
32+
int fd, int flags);
2833
TH_API char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx);
34+
TH_API int THMapAllocatorContext_fd(THMapAllocatorContext *ctx);
2935
TH_API long THMapAllocatorContext_size(THMapAllocatorContext *ctx);
3036
TH_API void THMapAllocatorContext_free(THMapAllocatorContext *ctx);
3137

0 commit comments

Comments
 (0)