Add zlib:set_controlling_process/2

5 files changed, 97 insertions, 18 deletions

diff --git a/erts/doc/src/zlib.xml b/erts/doc/src/zlib.xml
index f5cc1b1e64..6f4c42da27 100644
--- a/erts/doc/src/zlib.xml
+++ b/erts/doc/src/zlib.xml
@@ -71,6 +71,11 @@ list_to_binary([Compressed|Last])</pre>
       called prior to a call to
       <seealso marker="#inflate/2"><c>inflate/2</c></seealso>.
       </item>
+      <tag><c>not_on_controlling_process</c></tag>
+      <item>The stream was used by a process that doesn't control it. Use
+      <seealso marker="#set_controlling_process/2">
+      <c>set_controlling_process/2</c></seealso> if you need to transfer
+      a stream to a different process.</item>
       <tag><c>data_error</c></tag>
       <item>The data contains errors.
       </item>
@@ -740,6 +745,15 @@ loop(Z, Handler, {finished, Output}) ->
     </func>
 
     <func>
+      <name name="set_controlling_process" arity="2"/>
+      <fsummary>Transfers ownership of a zlib stream.</fsummary>
+      <desc>
+        <p>Changes the controlling process of <c><anno>Z</anno></c> to
+           <c><anno>Pid</anno></c>, which must be a local process.</p>
+      </desc>
+    </func>
+
+    <func>
       <name name="uncompress" arity="1"/>
       <fsummary>Uncompress data with standard zlib functionality.</fsummary>
       <desc>
diff --git a/erts/emulator/nifs/common/zlib_nif.c b/erts/emulator/nifs/common/zlib_nif.c
index a9c5b05e47..fa29b4fb71 100644
--- a/erts/emulator/nifs/common/zlib_nif.c
+++ b/erts/emulator/nifs/common/zlib_nif.c
@@ -106,6 +106,7 @@ typedef struct {
     int inflateChunk_buffer_size;
 
     ErlNifPid controlling_process;
+    ErlNifMutex *controller_lock;
 
     ErlNifIOQueue *input_queue;
 
@@ -117,6 +118,8 @@ typedef struct {
 
 static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
+static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
+
 static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_deflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 static ERL_NIF_TERM zlib_deflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
@@ -143,6 +146,11 @@ static ERL_NIF_TERM zlib_setBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM
 static ERL_NIF_TERM zlib_enqueue_input(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
 
 static ErlNifFunc nif_funcs[] = {
+    {"close_nif", 1, zlib_close},
+    {"open_nif", 0, zlib_open},
+
+    {"set_controller_nif", 2, zlib_set_controller},
+
     /* deflate */
     {"deflateInit_nif", 6, zlib_deflateInit},
     {"deflateSetDictionary_nif", 2, zlib_deflateSetDictionary},
@@ -162,10 +170,6 @@ static ErlNifFunc nif_funcs[] = {
     /* running checksum */
     {"crc32_nif", 1, zlib_crc32},
 
-    /* open & close */
-    {"close_nif", 1, zlib_close},
-    {"open_nif", 0, zlib_open},
-
     /* The stash keeps a single term alive across calls, and is used in
      * exception_on_need_dict/1 to retain the old error behavior, and for
      * saving data flushed through deflateParams/3. */
@@ -281,9 +285,7 @@ static ERL_NIF_TERM zlib_return(ErlNifEnv *env, int code) {
     return reason;
 }
 
-static void gc_zlib(ErlNifEnv *env, void* data) {
-    zlib_data_t *d = (zlib_data_t*)data;
-
+static void zlib_internal_close(zlib_data_t *d) {
     if(d->state == ST_DEFLATE) {
         deflateEnd(&d->s);
     } else if(d->state == ST_INFLATE) {
@@ -291,8 +293,6 @@ static void gc_zlib(ErlNifEnv *env, void* data) {
     }
 
     if(d->state != ST_CLOSED) {
-        enif_ioq_destroy(d->input_queue);
-
         if(d->stash_env != NULL) {
             enif_free_env(d->stash_env);
         }
@@ -301,17 +301,36 @@ static void gc_zlib(ErlNifEnv *env, void* data) {
     }
 }
 
+static void gc_zlib(ErlNifEnv *env, void* data) {
+    zlib_data_t *d = (zlib_data_t*)data;
+
+    enif_mutex_destroy(d->controller_lock);
+    enif_ioq_destroy(d->input_queue);
+
+    zlib_internal_close(d);
+
+    (void)env;
+}
+
 static int get_zlib_data(ErlNifEnv *env, ERL_NIF_TERM opaque, zlib_data_t **d) {
     return enif_get_resource(env, opaque, rtype_zlib, (void **)d);
 }
 
 static int zlib_process_check(ErlNifEnv *env, zlib_data_t *d) {
+    int is_controlling_process;
     ErlNifPid current_process;
 
     enif_self(env, &current_process);
 
-    return enif_is_identical(enif_make_pid(env, &current_process),
+    enif_mutex_lock(d->controller_lock);
+
+    is_controlling_process = enif_is_identical(
+        enif_make_pid(env, &current_process),
         enif_make_pid(env, &d->controlling_process));
+
+    enif_mutex_unlock(d->controller_lock);
+
+    return is_controlling_process;
 }
 
 static void zlib_reset_input(zlib_data_t *d) {
@@ -516,6 +535,8 @@ static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[
 
     d->input_queue = enif_ioq_create(ERL_NIF_IOQ_NORMAL);
 
+    d->controller_lock = enif_mutex_create("zlib_controller_lock");
+
     d->s.zalloc = zlib_alloc;
     d->s.zfree  = zlib_free;
     d->s.opaque = d;
@@ -556,7 +577,28 @@ static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv
         return enif_raise_exception(env, am_not_initialized);
     }
 
-    gc_zlib(env, d);
+    zlib_internal_close(d);
+
+    return am_ok;
+}
+
+static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
+    zlib_data_t *d;
+
+    ErlNifPid new_owner;
+
+    if(argc != 2 || !get_zlib_data(env, argv[0], &d)
+                 || !enif_get_local_pid(env, argv[1], &new_owner)) {
+        return enif_make_badarg(env);
+    } else if(!zlib_process_check(env, d)) {
+        return enif_raise_exception(env, am_not_on_controlling_process);
+    }
+
+    enif_mutex_lock(d->controller_lock);
+
+    d->controlling_process = new_owner;
+
+    enif_mutex_unlock(d->controller_lock);
 
     return am_ok;
 }
diff --git a/erts/preloaded/ebin/zlib.beam b/erts/preloaded/ebin/zlib.beam
index 5048bdb846..f388bc723a 100644
--- a/erts/preloaded/ebin/zlib.beam
+++ b/erts/preloaded/ebin/zlib.beam
diff --git a/erts/preloaded/src/zlib.erl b/erts/preloaded/src/zlib.erl
index 3170ab6351..03c9ae38a1 100644
--- a/erts/preloaded/src/zlib.erl
+++ b/erts/preloaded/src/zlib.erl
@@ -20,7 +20,8 @@
 
 -module(zlib).
 
--export([open/0,close/1,deflateInit/1,deflateInit/2,deflateInit/6,
+-export([open/0,close/1,set_controlling_process/2,
+         deflateInit/1,deflateInit/2,deflateInit/6,
          deflateSetDictionary/2,deflateReset/1,deflateParams/3,
          deflate/2,deflate/3,deflateEnd/1,
          inflateInit/1,inflateInit/2,inflateInit/3,
@@ -128,6 +129,14 @@ close(Z) ->
 close_nif(_Z) ->
     erlang:nif_error(undef).
 
+-spec set_controlling_process(Z, Pid) -> 'ok' when
+      Z :: zstream(),
+      Pid :: pid().
+set_controlling_process(Z, Pid) ->
+    set_controller_nif(Z, Pid).
+set_controller_nif(_Z, _Pid) ->
+    erlang:nif_error(undef).
+
 -spec deflateInit(Z) -> 'ok' when
       Z :: zstream().
 deflateInit(Z) ->
diff --git a/lib/kernel/test/zlib_SUITE.erl b/lib/kernel/test/zlib_SUITE.erl
index d17eded811..1afcd155b3 100644
--- a/lib/kernel/test/zlib_SUITE.erl
+++ b/lib/kernel/test/zlib_SUITE.erl
@@ -978,23 +978,37 @@ split_bin(Last,Acc) ->
 
 only_allow_owner(Config) when is_list(Config) ->
     Z = zlib:open(),
+    Owner = self(),
 
     ?m(ok, zlib:inflateInit(Z)),
     ?m(ok, zlib:inflateReset(Z)),
 
     {Pid, Ref} = spawn_monitor(
         fun() ->
-            ?m(?EXIT(not_on_controlling_process), zlib:inflateReset(Z))
+            ?m(?EXIT(not_on_controlling_process), zlib:inflateReset(Z)),
+            Owner ! '$transfer_ownership',
+            receive
+                '$ownership_transferred' ->
+                    ?m(ok, zlib:inflateReset(Z))
+            after 200 ->
+                ct:fail("Never received transfer signal.")
+            end
         end),
+    ownership_transfer_check(Z, Pid, Ref).
 
+ownership_transfer_check(Z, WorkerPid, Ref) ->
     receive
-        {'DOWN', Ref, process, Pid, _Reason} ->
-            ok
+        '$transfer_ownership' ->
+            zlib:set_controlling_process(Z, WorkerPid),
+            WorkerPid ! '$ownership_transferred',
+            ownership_transfer_check(Z, WorkerPid, Ref);
+        {'DOWN', Ref, process, WorkerPid, normal} ->
+            ok;
+        {'DOWN', Ref, process, WorkerPid, Reason} ->
+            ct:fail("Spawned worker crashed with reason ~p.", [Reason])
     after 200 ->
         ct:fail("Spawned worker timed out.")
-    end,
-
-    ?m(ok, zlib:inflateReset(Z)).
+    end.
 
 sub_heap_binaries(Config) when is_list(Config) ->
     Compressed = zlib:compress(<<"gurka">>),