7 files changed, 147 insertions, 90 deletions

diff --git a/src/keeper.c b/src/keeper.c
index ae3e2a8..7eda598 100644
--- a/src/keeper.c
+++ b/src/keeper.c
@@ -657,7 +657,7 @@ void init_keepers( Universe* U, lua_State* L)
         for( i = 0; i < nb_keepers; ++ i)                      // keepersUD
         {
                 // note that we will leak K if we raise an error later
-                lua_State* K = PROPAGATE_ALLOCF_ALLOC();
+                lua_State* K = create_state( U, L);
                 if( K == NULL)
                 {
                         (void) luaL_error( L, "init_keepers() failed while creating keeper states; out of memory");
diff --git a/src/lanes.c b/src/lanes.c
index c8e012c..ccb32c0 100644
--- a/src/lanes.c
+++ b/src/lanes.c
@@ -580,25 +580,6 @@ static bool_t selfdestruct_remove( Lane* s)
 }
 
 /*
-** mutex-protected allocator for use with Lua states that have non-threadsafe allocators (such as LuaJIT)
-*/
-struct ProtectedAllocator_s
-{
-        lua_Alloc allocF;
-        void* allocUD;
-        MUTEX_T lock;
-};
-void * protected_lua_Alloc( void *ud, void *ptr, size_t osize, size_t nsize)
-{
-        void* p;
-        struct ProtectedAllocator_s* s = (struct ProtectedAllocator_s*) ud;
-        MUTEX_LOCK( &s->lock);
-        p = s->allocF( s->allocUD, ptr, osize, nsize);
-        MUTEX_UNLOCK( &s->lock);
-        return p;
-}
-
-/*
 * Process end; cancel any still free-running threads
 */
 static int selfdestruct_gc( lua_State* L)
@@ -679,15 +660,9 @@ static int selfdestruct_gc( lua_State* L)
 
                 // If some lanes are currently cleaning after themselves, wait until they are done.
                 // They are no longer listed in the selfdestruct chain, but they still have to lua_close().
+                while( U->selfdestructing_count > 0)
                 {
-                        bool_t again = TRUE;
-                        do
-                        {
-                                MUTEX_LOCK( &U->selfdestruct_cs);
-                                again = (U->selfdestructing_count > 0) ? TRUE : FALSE;
-                                MUTEX_UNLOCK( &U->selfdestruct_cs);
-                                YIELD();
-                        } while( again);
+                        YIELD();
                 }
 
                 //---
@@ -727,6 +702,13 @@ static int selfdestruct_gc( lua_State* L)
                 }
         }
 
+        // If some lanes are currently cleaning after themselves, wait until they are done.
+        // They are no longer listed in the selfdestruct chain, but they still have to lua_close().
+        while( U->selfdestructing_count > 0)
+        {
+                YIELD();
+        }
+
         // necessary so that calling free_deep_prelude doesn't crash because linda_id expects a linda lightuserdata at absolute slot 1
         lua_settop( L, 0);
         // no need to mutex-protect this as all threads in the universe are gone at that point
@@ -740,18 +722,7 @@ static int selfdestruct_gc( lua_State* L)
         close_keepers( U, L);
 
         // remove the protected allocator, if any
-        {
-                void* ud;
-                lua_Alloc allocF = lua_getallocf( L, &ud);
-
-                if( allocF == protected_lua_Alloc)
-                {
-                        struct ProtectedAllocator_s* s = (struct ProtectedAllocator_s*) ud;
-                        lua_setallocf( L, s->allocF, s->allocUD);
-                        MUTEX_FREE( &s->lock);
-                        s->allocF( s->allocUD, s, sizeof( struct ProtectedAllocator_s), 0);
-                }
-        }
+        cleanup_allocator_function( U, L);
 
 #if HAVE_LANE_TRACKING
         MUTEX_FREE( &U->tracking_cs);
@@ -2097,24 +2068,6 @@ LUAG_FUNC( configure)
         DEBUGSPEW_CODE( fprintf( stderr, INDENT_BEGIN "%p: lanes.configure() BEGIN\n" INDENT_END, L));
         DEBUGSPEW_CODE( if( U) ++ U->debugspew_indent_depth);
 
-        lua_getfield( L, 1, "protect_allocator");                                            // settings protect_allocator
-        if( lua_toboolean( L, -1))
-        {
-                void* allocUD;
-                lua_Alloc allocF = lua_getallocf( L, &allocUD);
-                if( allocF != protected_lua_Alloc) // just in case
-                {
-                        struct ProtectedAllocator_s* s = (struct ProtectedAllocator_s*) allocF( allocUD, NULL, 0, sizeof( struct ProtectedAllocator_s));
-                        s->allocF = allocF;
-                        s->allocUD = allocUD;
-                        MUTEX_INIT( &s->lock);
-                        lua_setallocf( L, protected_lua_Alloc, s);
-                }
-        }
-        lua_pop( L, 1);                                                                      // settings
-        STACK_MID( L, 1);
-
-        // grab or create the universe
         if( U == NULL)
         {
                 U = universe_create( L);                                                           // settings universe
@@ -2144,6 +2097,7 @@ LUAG_FUNC( configure)
                 MUTEX_INIT( &U->deep_lock);
                 MUTEX_INIT( &U->mtid_lock);
                 U->selfdestruct_first = SELFDESTRUCT_END;
+                initialize_allocator_function( U, L);
                 initialize_on_state_create( U, L);
                 init_keepers( U, L);
                 STACK_MID( L, 1);
diff --git a/src/lanes.lua b/src/lanes.lua
index 6779095..15908fa 100644
--- a/src/lanes.lua
+++ b/src/lanes.lua
@@ -76,8 +76,7 @@ lanes.configure = function( settings_)
                 track_lanes = false,
                 demote_full_userdata = nil,
                 verbose_errors = false,
-                -- LuaJIT provides a thread-unsafe allocator by default, so we need to protect it when used in parallel lanes
-                protect_allocator = (package.loaded.jit and jit.version and package.loaded.ffi and (package.loaded.ffi.abi( "32bit") or package.loaded.ffi.abi( "gc64"))) and true or false
+                allocator = nil
         }
         local boolean_param_checker = function( val_)
                 -- non-'boolean-false' should be 'boolean-true' or nil
@@ -90,7 +89,10 @@ lanes.configure = function( settings_)
                         return type( val_) == "number" and val_ > 0
                 end,
                 with_timers = boolean_param_checker,
-                protect_allocator = boolean_param_checker,
+                allocator = function( val_)
+                        -- can be nil, "protected", or a function
+                        return val_ and (type( val_) == "function" or val_ == "protected") or true
+                end,
                 on_state_create = function( val_)
                         -- on_state_create may be nil or a function
                         return val_ and type( val_) == "function" or true
diff --git a/src/macros_and_utils.h b/src/macros_and_utils.h
index acbe690..e40a615 100644
--- a/src/macros_and_utils.h
+++ b/src/macros_and_utils.h
@@ -12,25 +12,7 @@
 #define inline __inline
 #endif
 
- // For some reason, LuaJIT 64bits doesn't support lua_newstate()
-#ifndef PROPAGATE_ALLOCF //you should #define PROPAGATE_ALLOCF 1 for LuaJIT in GC64 mode
-#if defined(LUA_JITLIBNAME) && (defined(__x86_64__) || defined(_M_X64))
- //#pragma message( "LuaJIT 64 bits detected: don't propagate allocf")
-#define PROPAGATE_ALLOCF 0
-#else // LuaJIT x64
- //#pragma message( "PUC-Lua detected: propagate allocf")
-#define PROPAGATE_ALLOCF 1
-#endif // LuaJIT x64
-#endif // PROPAGATE_ALLOCF defined
-#if PROPAGATE_ALLOCF
-#define PROPAGATE_ALLOCF_PREP( L) void* allocUD; lua_Alloc allocF = lua_getallocf( L, &allocUD)
-#define PROPAGATE_ALLOCF_ALLOC() lua_newstate( allocF, allocUD)
-#else // PROPAGATE_ALLOCF
-#define PROPAGATE_ALLOCF_PREP( L)
-#define PROPAGATE_ALLOCF_ALLOC() luaL_newstate()
-#endif // PROPAGATE_ALLOCF
-
-#define USE_DEBUG_SPEW 0
+ #define USE_DEBUG_SPEW 0
 #if USE_DEBUG_SPEW
 extern char const* debugspew_indent;
 #define INDENT_BEGIN "%.*s "
diff --git a/src/tools.c b/src/tools.c
index f8fc342..8885dea 100644
--- a/src/tools.c
+++ b/src/tools.c
@@ -105,6 +105,79 @@ void luaG_dump( lua_State* L)
 }
 #endif // _DEBUG
 
+// ################################################################################################
+
+static void* protected_lua_Alloc( void *ud, void *ptr, size_t osize, size_t nsize)
+{
+        void* p;
+        ProtectedAllocator* s = (ProtectedAllocator*) ud;
+        MUTEX_LOCK( &s->lock);
+        p = s->definition.allocF( s->definition.allocUD, ptr, osize, nsize);
+        MUTEX_UNLOCK( &s->lock);
+        return p;
+}
+
+static int luaG_provide_protected_allocator( lua_State* L)
+{
+        Universe* U = universe_get( L);
+        AllocatorDefinition* def = lua_newuserdata( L, sizeof(AllocatorDefinition));
+        def->allocF = protected_lua_Alloc;
+        def->allocUD = &U->protected_allocator;
+        return 1;
+}
+
+// Do I need to disable this when compiling for LuaJIT to prevent issues?
+void initialize_allocator_function( Universe* U, lua_State* L)
+{
+        STACK_CHECK( L, 0);
+        lua_getfield( L, -1, "allocator");                    // settings allocator|nil|"protected"
+        if( !lua_isnil( L, -1))
+        {
+                // store C function pointer in an internal variable
+                U->provide_allocator = lua_tocfunction( L, -1);     // settings allocator
+                if( U->provide_allocator != NULL)
+                {
+                        // make sure the function doesn't have upvalues
+                        char const* upname = lua_getupvalue( L, -1, 1);   // settings allocator upval?
+                        if( upname != NULL) // should be "" for C functions with upvalues if any
+                        {
+                                (void) luaL_error( L, "config.allocator() shouldn't have upvalues");
+                        }
+                        // remove this C function from the config table so that it doesn't cause problems
+                        // when we transfer the config table in newly created Lua states
+                        lua_pushnil( L);                                  // settings allocator nil
+                        lua_setfield( L, -3, "allocator");                // settings allocator
+                }
+                else if( lua_type( L, -1) == LUA_TSTRING)
+                {
+                        // initialize all we need for the protected allocator
+                        MUTEX_INIT( &U->protected_allocator.lock); // the mutex
+                        // and the original allocator to call from inside protection by the mutex
+                        U->protected_allocator.definition.allocF = lua_getallocf( L, &U->protected_allocator.definition.allocUD);
+                        // before a state is created, this function will be called to obtain the allocator
+                        U->provide_allocator = luaG_provide_protected_allocator;
+
+                        lua_setallocf( L, protected_lua_Alloc, &U->protected_allocator);
+                }
+        }
+        lua_pop( L, 1);                                       // settings
+        STACK_END( L, 0);
+}
+
+void cleanup_allocator_function( Universe* U, lua_State* L)
+{
+        // remove the protected allocator, if any
+        if( U->protected_allocator.definition.allocF != NULL)
+        {
+                // install the non-protected allocator
+                lua_setallocf( L, U->protected_allocator.definition.allocF, U->protected_allocator.definition.allocUD);
+                // release the mutex
+                MUTEX_FREE( &U->protected_allocator.lock);
+        }
+}
+
+// ################################################################################################
+
 void initialize_on_state_create( Universe* U, lua_State* L)
 {
         STACK_CHECK( L, 0);
@@ -629,6 +702,31 @@ void call_on_state_create( Universe* U, lua_State* L, lua_State* from_, LookupMo
         }
 }
 
+lua_State* create_state( Universe* U, lua_State* from_)
+{
+        lua_State* L;
+        if( U->provide_allocator != NULL)
+        {
+                lua_pushcclosure( from_, U->provide_allocator, 0);
+                lua_call( from_, 0, 1);
+                {
+                        AllocatorDefinition* def = lua_touserdata( from_, -1);
+                        L = lua_newstate( def->allocF, def->allocUD);
+                }
+                lua_pop( from_, 1);
+        }
+        else
+        {
+                L = luaL_newstate();
+        }
+
+        if( L == NULL)
+        {
+                (void) luaL_error( from_, "luaG_newstate() failed while creating state; out of memory");
+        }
+        return L;
+}
+
 /*
  * Like 'luaL_openlibs()' but allows the set of libraries be selected
  *
@@ -644,16 +742,7 @@ void call_on_state_create( Universe* U, lua_State* L, lua_State* from_, LookupMo
  */
 lua_State* luaG_newstate( Universe* U, lua_State* from_, char const* libs_)
 {
-        // re-use alloc function from the originating state
-#if PROPAGATE_ALLOCF
-        PROPAGATE_ALLOCF_PREP( from_);
-#endif // PROPAGATE_ALLOCF
-        lua_State* L = PROPAGATE_ALLOCF_ALLOC();
-
-        if( L == NULL)
-        {
-                (void) luaL_error( from_, "luaG_newstate() failed while creating state; out of memory");
-        }
+        lua_State* L = create_state( U, from_);
 
         STACK_GROW( L, 2);
         STACK_CHECK_ABS( L, 0);
diff --git a/src/tools.h b/src/tools.h
index ca0b9fc..71460c3 100644
--- a/src/tools.h
+++ b/src/tools.h
@@ -22,6 +22,7 @@ typedef struct s_Universe Universe;
 void luaG_dump( lua_State* L);
 #endif // _DEBUG
 
+lua_State* create_state( Universe* U, lua_State* from_);
 lua_State* luaG_newstate( Universe* U, lua_State* _from, char const* libs);
 
 // ################################################################################################
@@ -36,6 +37,9 @@ int luaG_new_require( lua_State* L);
 
 void populate_func_lookup_table( lua_State* L, int _i, char const* _name);
 void serialize_require( DEBUGSPEW_PARAM_COMMA( Universe* U) lua_State *L);
+void initialize_allocator_function( Universe* U, lua_State* L);
+void cleanup_allocator_function( Universe* U, lua_State* L);
+
 void initialize_on_state_create( Universe* U, lua_State* L);
 void call_on_state_create( Universe* U, lua_State* L, lua_State* from_, LookupMode mode_);
 
diff --git a/src/universe.h b/src/universe.h
index 359dc90..8727bf7 100644
--- a/src/universe.h
+++ b/src/universe.h
@@ -24,6 +24,24 @@ typedef struct s_Lane Lane;
 
 // ################################################################################################
 
+// everything we need to provide to lua_newstate()
+struct AllocatorDefinition_s
+{
+        lua_Alloc allocF;
+        void* allocUD;
+};
+typedef struct AllocatorDefinition_s AllocatorDefinition;
+
+// mutex-protected allocator for use with Lua states that share a non-threadsafe allocator
+struct ProtectedAllocator_s
+{
+        AllocatorDefinition definition;
+        MUTEX_T lock;
+};
+typedef struct ProtectedAllocator_s ProtectedAllocator;
+
+// ################################################################################################
+
 // everything regarding the a Lanes universe is stored in that global structure
 // held as a full userdata in the master Lua state that required it for the first time
 // don't forget to initialize all members in LG_configure()
@@ -34,8 +52,16 @@ struct s_Universe
 
         bool_t demoteFullUserdata;
 
+        // before a state is created, this function will be called to obtain the allocator
+        lua_CFunction provide_allocator;
+
+        // after a state is created, this function will be called right after the bases libraries are loaded
         lua_CFunction on_state_create_func;
 
+        // Initialized and used only if allocator="protected" is found in the configuration settings
+        // contains a mutex and the original allocator definition
+        ProtectedAllocator protected_allocator;
+
         Keepers* keepers;
 
         // Initialized by 'init_once_LOCKED()': the deep userdata Linda object