erts: Add enif_select_read|write with 'msg_env' argument

7 files changed, 122 insertions, 64 deletions

diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 7339aa8874..a48d0391f6 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -707,6 +707,29 @@ error:
     return reds;
 }
 
+/** @brief Create a message with the content of process independent \c msg_env.
+ *  Invalidates \c msg_env.
+ */
+ErtsMessage* erts_create_message_from_nif_env(ErlNifEnv* msg_env)
+{
+    struct enif_msg_environment_t* menv = (struct enif_msg_environment_t*)msg_env;
+    ErtsMessage* mp;
+
+    flush_env(msg_env);
+    mp = erts_alloc_message(0, NULL);
+    mp->data.heap_frag = menv->env.heap_frag;
+    ASSERT(mp->data.heap_frag == MBUF(&menv->phony_proc));
+    if (mp->data.heap_frag != NULL) {
+        /* Move all offheap's from phony proc to the first fragment.
+           Quick and dirty... */
+        ASSERT(!is_offheap(&mp->data.heap_frag->off_heap));
+        mp->data.heap_frag->off_heap = MSO(&menv->phony_proc);
+        clear_offheap(&MSO(&menv->phony_proc));
+        menv->env.heap_frag = NULL;
+        MBUF(&menv->phony_proc) = NULL;
+    }
+    return mp;
+}
 
 int enif_send(ErlNifEnv* env, const ErlNifPid* to_pid,
 	      ErlNifEnv* msg_env, ERL_NIF_TERM msg)
@@ -803,20 +826,8 @@ int enif_send(ErlNifEnv* env, const ErlNifPid* to_pid,
             }
 #endif
         }
-        flush_env(msg_env);
-        mp = erts_alloc_message(0, NULL);
+        mp = erts_create_message_from_nif_env(msg_env);
         ERL_MESSAGE_TOKEN(mp) = token;
-        mp->data.heap_frag = menv->env.heap_frag;
-        ASSERT(mp->data.heap_frag == MBUF(&menv->phony_proc));
-        if (mp->data.heap_frag != NULL) {
-            /* Move all offheap's from phony proc to the first fragment.
-               Quick and dirty... */
-            ASSERT(!is_offheap(&mp->data.heap_frag->off_heap));
-            mp->data.heap_frag->off_heap = MSO(&menv->phony_proc);
-            clear_offheap(&MSO(&menv->phony_proc));
-            menv->env.heap_frag = NULL;
-            MBUF(&menv->phony_proc) = NULL;
-        }
     } else {
         erts_literal_area_t litarea;
 	ErlOffHeap *ohp;
diff --git a/erts/emulator/beam/erl_nif.h b/erts/emulator/beam/erl_nif.h
index 58a217c20b..3fd1a8fd4c 100644
--- a/erts/emulator/beam/erl_nif.h
+++ b/erts/emulator/beam/erl_nif.h
@@ -54,7 +54,7 @@
 ** 2.13: 20.1 add enif_ioq
 ** 2.14: 21.0 add enif_ioq_peek_head, enif_(mutex|cond|rwlock|thread)_name
 **                enif_vfprintf, enif_vsnprintf, enif_make_map_from_arrays
-** 2.15: 22.0 ERL_NIF_SELECT_CANCEL
+** 2.15: 22.0 ERL_NIF_SELECT_CANCEL, enif_select_(read|write)
 */
 #define ERL_NIF_MAJOR_VERSION 2
 #define ERL_NIF_MINOR_VERSION 15
diff --git a/erts/emulator/beam/erl_nif_api_funcs.h b/erts/emulator/beam/erl_nif_api_funcs.h
index 81f64f2390..129166562d 100644
--- a/erts/emulator/beam/erl_nif_api_funcs.h
+++ b/erts/emulator/beam/erl_nif_api_funcs.h
@@ -210,6 +210,9 @@ ERL_NIF_API_FUNC_DECL(int,enif_vsnprintf,(char*, size_t, const char *fmt, va_lis
 
 ERL_NIF_API_FUNC_DECL(int,enif_make_map_from_arrays,(ErlNifEnv *env, ERL_NIF_TERM keys[], ERL_NIF_TERM values[], size_t cnt, ERL_NIF_TERM *map_out));
 
+ERL_NIF_API_FUNC_DECL(int,enif_select_x,(ErlNifEnv* env, ErlNifEvent e, enum ErlNifSelectFlags flags, void* obj, const ErlNifPid* pid, ERL_NIF_TERM msg, ErlNifEnv* msg_env));
+
+
 /*
 ** ADD NEW ENTRIES HERE (before this comment) !!!
 */
@@ -392,6 +395,7 @@ ERL_NIF_API_FUNC_DECL(int,enif_make_map_from_arrays,(ErlNifEnv *env, ERL_NIF_TER
 #  define enif_vfprintf ERL_NIF_API_FUNC_MACRO(enif_vfprintf)
 #  define enif_vsnprintf ERL_NIF_API_FUNC_MACRO(enif_vsnprintf)
 #  define enif_make_map_from_arrays ERL_NIF_API_FUNC_MACRO(enif_make_map_from_arrays)
+#  define enif_select_x ERL_NIF_API_FUNC_MACRO(enif_select_x)
 
 /*
 ** ADD NEW ENTRIES HERE (before this comment)
@@ -623,6 +627,12 @@ static ERL_NIF_INLINE ERL_NIF_TERM enif_make_list9(ErlNifEnv* env,
 #ifndef enif_make_pid
 
 #  define enif_make_pid(ENV, PID) ((void)(ENV),(const ERL_NIF_TERM)((PID)->pid))
+#  define enif_select_read(ENV, E, OBJ, PID, MSG, MSG_ENV) \
+    enif_select_x(ENV, E, ERL_NIF_SELECT_READ | ERL_NIF_SELECT_CUSTOM_MSG, \
+                  OBJ, PID, MSG, MSG_ENV)
+#  define enif_select_write(ENV, E, OBJ, PID, MSG, MSG_ENV) \
+    enif_select_x(ENV, E, ERL_NIF_SELECT_WRITE | ERL_NIF_SELECT_CUSTOM_MSG, \
+                  OBJ, PID, MSG, MSG_ENV)
 
 #if SIZEOF_LONG == 8
 #  define enif_get_int64 enif_get_long
diff --git a/erts/emulator/beam/global.h b/erts/emulator/beam/global.h
index 77b5a3ca05..9eb7b58dbb 100644
--- a/erts/emulator/beam/global.h
+++ b/erts/emulator/beam/global.h
@@ -131,6 +131,7 @@ extern Eterm erts_nif_call_function(Process *p, Process *tracee,
 
 int erts_call_dirty_nif(ErtsSchedulerData *esdp, Process *c_p,
 			BeamInstr *I, Eterm *reg);
+ErtsMessage* erts_create_message_from_nif_env(ErlNifEnv* msg_env);
 
 
 /* Driver handle (wrapper for old plain handle) */
diff --git a/erts/emulator/sys/common/erl_check_io.c b/erts/emulator/sys/common/erl_check_io.c
index 487b9cba40..304e0a5d0c 100644
--- a/erts/emulator/sys/common/erl_check_io.c
+++ b/erts/emulator/sys/common/erl_check_io.c
@@ -596,11 +596,11 @@ static void prepare_select_msg(struct erts_nif_select_event* e,
                                Eterm recipient,
                                ErtsResource* resource,
                                Eterm msg,
+                               ErlNifEnv* msg_env,
                                Eterm event_atom)
 {
     ErtsMessage* mp;
     Eterm* hp;
-    Eterm* hp_start;
     Uint hsz;
 
     if (is_not_nil(e->pid)) {
@@ -609,14 +609,20 @@ static void prepare_select_msg(struct erts_nif_select_event* e,
     }
 
     if (mode & ERL_NIF_SELECT_CUSTOM_MSG) {
-        hsz = size_object(msg);
-        mp = erts_alloc_message(hsz, &hp);
-        hp_start = hp;
-        ERL_MESSAGE_TERM(mp) = copy_struct(msg, hsz, &hp, &mp->hfrag.off_heap);
+        if (msg_env) {
+            mp = erts_create_message_from_nif_env(msg_env);
+            ERL_MESSAGE_TERM(mp) = msg;
+        }
+        else {
+            hsz = size_object(msg);
+            mp = erts_alloc_message(hsz, &hp);
+            ERL_MESSAGE_TERM(mp) = copy_struct(msg, hsz, &hp, &mp->hfrag.off_heap);
+        }
     }
     else {
         ErtsBinary* bin;
         Eterm resource_term, ref_term, tuple;
+        Eterm* hp_start;
 
          /* {select, Resource, Ref, EventAtom} */
         hsz = 5 + ERTS_MAGIC_REF_THING_SIZE;
@@ -643,8 +649,8 @@ static void prepare_select_msg(struct erts_nif_select_event* e,
         tuple = TUPLE4(hp, am_select, resource_term, ref_term, event_atom);
         hp += 5;
         ERL_MESSAGE_TERM(mp) = tuple;
+        ASSERT(hp == hp_start + hsz); (void)hp_start;
     }
-    ASSERT(hp == hp_start + hsz); (void)hp_start;
 
     ASSERT(is_not_nil(recipient));
     e->pid = recipient;
@@ -1027,12 +1033,21 @@ done_unknown:
 }
 
 int
-enif_select(ErlNifEnv* env,
-            ErlNifEvent e,
-            enum ErlNifSelectFlags mode,
-            void* obj,
-            const ErlNifPid* pid,
-            Eterm msg)
+enif_select(ErlNifEnv* env, ErlNifEvent e, enum ErlNifSelectFlags mode,
+            void* obj, const ErlNifPid* pid, Eterm msg)
+{
+    return enif_select_x(env, e, mode, obj, pid, msg, NULL);
+}
+
+
+int
+enif_select_x(ErlNifEnv* env,
+              ErlNifEvent e,
+              enum ErlNifSelectFlags mode,
+              void* obj,
+              const ErlNifPid* pid,
+              Eterm msg,
+              ErlNifEnv* msg_env)
 {
     int on;
     ErtsResource* resource = DATA_TO_RESOURCE(obj);
@@ -1177,11 +1192,12 @@ enif_select(ErlNifEnv* env,
         ASSERT(state->driver.stop.resource == resource);
         if (mode & ERL_DRV_READ) {
             prepare_select_msg(&state->driver.nif->in, mode, recipient,
-                               resource, msg, am_ready_input);
+                               resource, msg, msg_env, am_ready_input);
+            msg_env = NULL;
         }
         if (mode & ERL_DRV_WRITE) {
             prepare_select_msg(&state->driver.nif->out, mode, recipient,
-                               resource, msg, am_ready_output);
+                               resource, msg, msg_env, am_ready_output);
         }
         ret = 0;
     }
diff --git a/erts/emulator/test/nif_SUITE.erl b/erts/emulator/test/nif_SUITE.erl
index 8de6d256c5..0d0930e124 100644
--- a/erts/emulator/test/nif_SUITE.erl
+++ b/erts/emulator/test/nif_SUITE.erl
@@ -497,79 +497,82 @@ t_on_load(Config) when is_list(Config) ->
 
 select(Config) when is_list(Config) ->
     ensure_lib_loaded(Config),
-    select_do(0, make_ref(), make_ref()),
+    select_do(0, make_ref(), make_ref(), null),
 
     RefBin = list_to_binary(lists:duplicate(100, $x)),
-    select_do(?ERL_NIF_SELECT_CUSTOM_MSG,
-              small, {a, tuple, with, "some", RefBin}),
+    [select_do(?ERL_NIF_SELECT_CUSTOM_MSG,
+               small, {a, tuple, with, "some", RefBin}, MSG_ENV)
+     || MSG_ENV <- [null, alloc_env]],
     ok.
 
-select_do(Flag, Ref, Ref2) ->
+select_do(Flag, Ref, Ref2, MSG_ENV) ->
+    io:format("select_do(~p, ~p, ~p)\n", [Ref, Ref2, MSG_ENV]),
+
     {{R, R_ptr}, {W, W_ptr}} = pipe_nif(),
     ok = write_nif(W, <<"hej">>),
     <<"hej">> = read_nif(R, 3),
 
     %% Wait for read
     eagain = read_nif(R, 3),
-    0 = select_nif(R,?ERL_NIF_SELECT_READ bor Flag, R,null,Ref),
+    0 = select_nif(R,?ERL_NIF_SELECT_READ bor Flag, R,null,Ref,MSG_ENV),
     [] = flush(0),
     ok = write_nif(W, <<"hej">>),
     receive_ready(R, Ref, ready_input),
-    0 = select_nif(R,?ERL_NIF_SELECT_READ bor Flag,R,self(),Ref2),
+    0 = select_nif(R,?ERL_NIF_SELECT_READ bor Flag,R,self(),Ref2,MSG_ENV),
     receive_ready(R, Ref2, ready_input),
     Papa = self(),
     Pid = spawn_link(fun() ->
                              receive_ready(R, Ref, ready_input),
                              Papa ! {self(), done}
                      end),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, Pid, Ref),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, Pid, Ref,MSG_ENV),
     {Pid, done} = receive_any(1000),
 
     %% Cancel read
-    0 = select_nif(R,?ERL_NIF_SELECT_READ bor ?ERL_NIF_SELECT_CANCEL,R,null,Ref),
+    0 = select_nif(R,?ERL_NIF_SELECT_READ bor ?ERL_NIF_SELECT_CANCEL,R,null,Ref,null),
     <<"hej">> = read_nif(R, 3),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref, MSG_ENV),
     ?ERL_NIF_SELECT_READ_CANCELLED =
-        select_nif(R,?ERL_NIF_SELECT_READ bor ?ERL_NIF_SELECT_CANCEL,R,null,Ref),
+        select_nif(R,?ERL_NIF_SELECT_READ bor ?ERL_NIF_SELECT_CANCEL,R,null,Ref,null),
     ok = write_nif(W, <<"hej again">>),
     [] = flush(0),
     <<"hej again">> = read_nif(R, 9),
 
     %% Wait for write
     Written = write_full(W, $a),
-    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor Flag, W, self(), Ref),
+    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor Flag, W, self(), Ref, MSG_ENV),
     [] = flush(0),
     Written = read_nif(R,byte_size(Written)),
     receive_ready(W, Ref, ready_output),
 
     %% Cancel write
-    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor ?ERL_NIF_SELECT_CANCEL, W, null, Ref),
+    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor ?ERL_NIF_SELECT_CANCEL, W, null, Ref, null),
     Written2 = write_full(W, $b),
-    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor Flag, W, null, Ref),
+    0 = select_nif(W, ?ERL_NIF_SELECT_WRITE bor Flag, W, null, Ref, MSG_ENV),
     ?ERL_NIF_SELECT_WRITE_CANCELLED =
-        select_nif(W, ?ERL_NIF_SELECT_WRITE bor ?ERL_NIF_SELECT_CANCEL, W, null, Ref),
+        select_nif(W, ?ERL_NIF_SELECT_WRITE bor ?ERL_NIF_SELECT_CANCEL, W, null, Ref, null),
     Written2 = read_nif(R,byte_size(Written2)),
     [] = flush(0),
 
     %% Close write and wait for EOF
     eagain = read_nif(R, 1),
-    check_stop_ret(select_nif(W, ?ERL_NIF_SELECT_STOP, W, null, Ref)),
+    check_stop_ret(select_nif(W, ?ERL_NIF_SELECT_STOP, W, null, Ref, null)),
     [{fd_resource_stop, W_ptr, _}] = flush(),
     {1, {W_ptr,_}} = last_fd_stop_call(),
     true = is_closed_nif(W),
     [] = flush(0),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, self(), Ref),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, self(), Ref, MSG_ENV),
     receive_ready(R, Ref, ready_input),
     eof = read_nif(R,1),
 
-    check_stop_ret(select_nif(R, ?ERL_NIF_SELECT_STOP, R, null, Ref)),
+    check_stop_ret(select_nif(R, ?ERL_NIF_SELECT_STOP, R, null, Ref, null)),
     [{fd_resource_stop, R_ptr, _}] = flush(),
     {1, {R_ptr,_}} = last_fd_stop_call(),
     true = is_closed_nif(R),
 
-    select_2(Flag, Ref, Ref2).
+    select_2(Flag, Ref, Ref2, MSG_ENV).
 
-select_2(Flag, Ref1, Ref2) ->
+select_2(Flag, Ref1, Ref2, MSG_ENV) ->
     erlang:garbage_collect(),
     {_,_,2} = last_resource_dtor_call(),
 
@@ -577,8 +580,8 @@ select_2(Flag, Ref1, Ref2) ->
 
     %% Change ref
     eagain = read_nif(R, 1),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, self(), Ref2),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1, MSG_ENV),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, self(), Ref2, MSG_ENV),
 
     [] = flush(0),
     ok = write_nif(W, <<"hej">>),
@@ -587,10 +590,10 @@ select_2(Flag, Ref1, Ref2) ->
 
     %% Change pid
     eagain = read_nif(R, 1),
-    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1),
+    0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1, MSG_ENV),
     Papa = self(),
     spawn_link(fun() ->
-                       0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1),
+                       0 = select_nif(R, ?ERL_NIF_SELECT_READ bor Flag, R, null, Ref1, MSG_ENV),
                        [] = flush(0),
                        Papa ! sync,
                        receive_ready(R, Ref1, ready_input),
@@ -602,13 +605,13 @@ select_2(Flag, Ref1, Ref2) ->
     done = receive_any(),
     [] = flush(0),
 
-    check_stop_ret(select_nif(R,?ERL_NIF_SELECT_STOP,R,null,Ref1)),
+    check_stop_ret(select_nif(R,?ERL_NIF_SELECT_STOP,R,null,Ref1, null)),
     [{fd_resource_stop, R_ptr, _}] = flush(),
     {1, {R_ptr,_}} = last_fd_stop_call(),
     true = is_closed_nif(R),
 
     %% Stop without previous read/write select
-    ?ERL_NIF_SELECT_STOP_CALLED = select_nif(W,?ERL_NIF_SELECT_STOP,W,null,Ref1),
+    ?ERL_NIF_SELECT_STOP_CALLED = select_nif(W,?ERL_NIF_SELECT_STOP,W,null,Ref1,null),
     [{fd_resource_stop, W_ptr, 1}] = flush(),
     {1, {W_ptr,1}} = last_fd_stop_call(),
     true = is_closed_nif(W),
@@ -634,7 +637,7 @@ select_steal_child_process(Parent, RFd) ->
     Ref2 = make_ref(),
 
     %% Try to select from the child pid (steal from parent)
-    ?assertEqual(0, select_nif(R2Fd, ?ERL_NIF_SELECT_READ, R2Fd, null, Ref2)),
+    ?assertEqual(0, select_nif(R2Fd, ?ERL_NIF_SELECT_READ, R2Fd, null, Ref2, null)),
     ?assertEqual([], flush(0)),
     ?assertEqual(eagain, read_nif(R2Fd, 1)),
 
@@ -642,7 +645,7 @@ select_steal_child_process(Parent, RFd) ->
     Parent ! {self(), stage1}, % signal parent to send the <<"stolen1">>
 
     %% Receive <<"stolen1">> via enif_select
-    ?assertEqual(0, select_nif(R2Fd, ?ERL_NIF_SELECT_READ, R2Fd, null, Ref2)),
+    ?assertEqual(0, select_nif(R2Fd, ?ERL_NIF_SELECT_READ, R2Fd, null, Ref2, null)),
     ?assertMatch([{select, R2Fd, Ref2, ready_input}], flush()),
     ?assertEqual(<<"stolen1">>, read_nif(R2Fd, 7)),
 
@@ -660,7 +663,7 @@ select_steal(Config) when is_list(Config) ->
     {{RFd, RPtr}, {WFd, WPtr}} = pipe_nif(),
 
     %% Bind the socket to current pid in enif_select
-    ?assertEqual(0, select_nif(RFd, ?ERL_NIF_SELECT_READ, RFd, null, Ref)),
+    ?assertEqual(0, select_nif(RFd, ?ERL_NIF_SELECT_READ, RFd, null, Ref, null)),
     ?assertEqual([], flush(0)),
 
     %% Spawn a process and do some stealing
@@ -674,15 +677,15 @@ select_steal(Config) when is_list(Config) ->
     ?assertMatch([{Pid, done}], flush(1)),  % synchronize with the child
 
     %% Try to select from the parent pid (steal back)
-    ?assertEqual(0, select_nif(RFd, ?ERL_NIF_SELECT_READ, RFd, Pid, Ref)),
+    ?assertEqual(0, select_nif(RFd, ?ERL_NIF_SELECT_READ, RFd, Pid, Ref, null)),
 
     %% Ensure that no data is hanging and close.
     %% Rfd is stolen at this point.
-    check_stop_ret(select_nif(WFd, ?ERL_NIF_SELECT_STOP, WFd, null, Ref)),
+    check_stop_ret(select_nif(WFd, ?ERL_NIF_SELECT_STOP, WFd, null, Ref, null)),
     ?assertMatch([{fd_resource_stop, WPtr, _}], flush()),
     {1, {WPtr, 1}} = last_fd_stop_call(),
 
-    check_stop_ret(select_nif(RFd, ?ERL_NIF_SELECT_STOP, RFd, null, Ref)),
+    check_stop_ret(select_nif(RFd, ?ERL_NIF_SELECT_STOP, RFd, null, Ref, null)),
     ?assertMatch([{fd_resource_stop, RPtr, _}], flush()),
     {1, {RPtr, _DirectCall}} = last_fd_stop_call(),
 
@@ -3406,7 +3409,7 @@ term_to_binary_nif(_, _) -> ?nif_stub.
 binary_to_term_nif(_, _, _) -> ?nif_stub.
 port_command_nif(_, _) -> ?nif_stub.
 format_term_nif(_,_) -> ?nif_stub.
-select_nif(_,_,_,_,_) -> ?nif_stub.
+select_nif(_,_,_,_,_,_) -> ?nif_stub.
 dupe_resource_nif(_) -> ?nif_stub.
 pipe_nif() -> ?nif_stub.
 write_nif(_,_) -> ?nif_stub.
diff --git a/erts/emulator/test/nif_SUITE_data/nif_SUITE.c b/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
index f2ce6dbe67..21af4b05b3 100644
--- a/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
+++ b/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
@@ -2486,7 +2486,8 @@ static ERL_NIF_TERM select_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv
     enum ErlNifSelectFlags mode;
     void* obj;
     ErlNifPid nifpid, *pid = NULL;
-    ERL_NIF_TERM ref;
+    ERL_NIF_TERM ref_or_msg;
+    ErlNifEnv* msg_env = NULL;
     int retval;
 
     if (!get_fd(env, argv[0], &fdr)
@@ -2501,11 +2502,27 @@ static ERL_NIF_TERM select_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv
 	    return enif_make_badarg(env);
 	pid = &nifpid;
     }
-    ref = argv[4];
+    ref_or_msg = argv[4];
+    if (argv[5] != atom_null) {
+        msg_env = enif_alloc_env();
+        ref_or_msg = enif_make_copy(msg_env, ref_or_msg);
+    }
 
     fdr->was_selected = 1;
     enif_self(env, &fdr->pid);
-    retval = enif_select(env, fdr->fd, mode, obj, pid, ref);
+    switch (mode) {
+    case ERL_NIF_SELECT_CUSTOM_MSG | ERL_NIF_SELECT_READ:
+        retval = enif_select_read(env, fdr->fd, obj, pid, ref_or_msg, msg_env);
+        break;
+    case ERL_NIF_SELECT_CUSTOM_MSG | ERL_NIF_SELECT_WRITE:
+        retval = enif_select_write(env, fdr->fd, obj, pid, ref_or_msg, msg_env);
+        break;
+    default:
+        retval = enif_select(env, fdr->fd, mode, obj, pid, ref_or_msg);
+    }
+
+    if (msg_env)
+        enif_free_env(msg_env);
 
     return enif_make_int(env, retval);
 }
@@ -3565,7 +3582,7 @@ static ErlNifFunc nif_funcs[] =
     {"binary_to_term_nif", 3, binary_to_term},
     {"port_command_nif", 2, port_command},
     {"format_term_nif", 2, format_term},
-    {"select_nif", 5, select_nif},
+    {"select_nif", 6, select_nif},
 #ifndef __WIN32__
     {"pipe_nif", 0, pipe_nif},
     {"write_nif", 2, write_nif},