@@ -60,26 +60,32 @@ def __setstate__(self, state):
60
60
else :
61
61
62
62
class Arena (object ):
63
+ if sys .platform == 'linux' :
64
+ _dir_candidates = ['/dev/shm' ]
65
+ else :
66
+ _dir_candidates = []
63
67
64
68
def __init__ (self , size , fd = - 1 ):
65
69
self .size = size
66
70
self .fd = fd
67
71
if fd == - 1 :
68
72
self .fd , name = tempfile .mkstemp (
69
- prefix = 'pym-%d-' % os .getpid (), dir = util .get_temp_dir ())
73
+ prefix = 'pym-%d-' % os .getpid (),
74
+ dir = self ._choose_dir (size ))
70
75
os .unlink (name )
71
76
util .Finalize (self , os .close , (self .fd ,))
72
- with open (self .fd , 'wb' , closefd = False ) as f :
73
- bs = 1024 * 1024
74
- if size >= bs :
75
- zeros = b'\0 ' * bs
76
- for _ in range (size // bs ):
77
- f .write (zeros )
78
- del zeros
79
- f .write (b'\0 ' * (size % bs ))
80
- assert f .tell () == size
77
+ os .ftruncate (self .fd , size )
81
78
self .buffer = mmap .mmap (self .fd , self .size )
82
79
80
+ def _choose_dir (self , size ):
81
+ # Choose a non-storage backed directory if possible,
82
+ # to improve performance
83
+ for d in self ._dir_candidates :
84
+ st = os .statvfs (d )
85
+ if st .f_bavail * st .f_frsize >= size : # enough free space?
86
+ return d
87
+ return util .get_temp_dir ()
88
+
83
89
def reduce_arena (a ):
84
90
if a .fd == - 1 :
85
91
raise ValueError ('Arena is unpicklable because '
0 commit comments