Add possibility to also get size of data from erlang:dist_ctrl_get_data()

10 files changed, 225 insertions, 5 deletions

diff --git a/erts/doc/src/erlang.xml b/erts/doc/src/erlang.xml
index e78ded4ae1..0d94f83493 100644
--- a/erts/doc/src/erlang.xml
+++ b/erts/doc/src/erlang.xml
@@ -1269,6 +1269,42 @@ end</code>
 	  data is available by calling
 	  <seealso marker="erlang#dist_ctrl_get_data_notification/1"><c>erlang:dist_ctrl_get_data_notification(DHandle)</c></seealso>.
 	</p>
+	<p>The returned value when there are data available depends
+	on the value of the <c>get_size</c> option configured on the
+	distribution channel identified by <c><anno>DHandle</anno></c>.
+	For more information see the documentation of the <c>get_size</c>
+	option for the
+	<seealso marker="#dist_ctrl_set_opt/3"><c>erlang:dist_ctrl_set_opt/3</c></seealso>
+	function.</p>
+	<note><p>
+	  Only the process registered as distribution
+	  controller for the distribution channel identified by
+	  <c><anno>DHandle</anno></c> is allowed to call this
+	  function.
+	</p></note>
+	<p>
+	  This function is used when implementing an alternative
+	  distribution carrier using processes as distribution
+	  controllers. <c><anno>DHandle</anno></c> is retrived
+	  via the callback
+	  <seealso marker="erts:alt_dist#hs_data_f_handshake_complete"><c>f_handshake_complete</c></seealso>.
+	  More information can be found in the documentation of
+	  <seealso marker="erts:alt_dist#distribution_module">ERTS
+	  User's Guide ➜ How to implement an Alternative Carrier
+	  for the Erlang Distribution ➜ Distribution Module</seealso>.
+	</p>
+      </desc>
+    </func>
+
+    <func>
+      <name name="dist_ctrl_get_opt" arity="2" clause_i="1" since="OTP @OTP-15617@"/>
+      <fsummary>Get value of the get_size option on a distribution channel</fsummary>
+      <desc>
+	<p>Returns the value of the <c>get_size</c> option on the distribution channel
+	identified by <c><anno>DHandle</anno></c>. For more information see the
+	documentation of the <c>get_size</c> option for the
+	<seealso marker="#dist_ctrl_set_opt/3"><c>erlang:dist_ctrl_set_opt/3</c></seealso>
+	function.</p>
 	<note><p>
 	  Only the process registered as distribution
 	  controller for the distribution channel identified by
@@ -1392,6 +1428,55 @@ end</code>
     </func>
 
     <func>
+      <name name="dist_ctrl_set_opt" arity="3" clause_i="1" since="OTP @OTP-15617@"/>
+      <fsummary>Set value of the get_size option on a distribution channel</fsummary>
+      <desc>
+        <p>Sets the value of the <c>get_size</c> option on the distribution channel
+	identified by <c><anno>DHandle</anno></c>. This option controls the return
+	value of calls to
+	<seealso marker="#dist_ctrl_get_data/1">erlang:dist_ctrl_get_data(<anno>DHandle</anno>)</seealso>
+	where <c><anno>DHandle</anno></c> equals <c><anno>DHandle</anno></c> used
+	when setting this option.
+	When the <c>get_size</c> option is:</p>
+	<taglist>
+	  <tag><c>false</c></tag>
+	  <item>
+	    and there are distribution data available, a call to
+	    <c>erlang:dist_ctrl_get_data(<anno>DHandle</anno>)</c>
+	    will just return <c>Data</c> to pass over the channel.
+	    This is the default value of the <c>get_size</c> option.
+	  </item>
+	  <tag><c>true</c></tag>
+	  <item>
+	    and there are distribution data available, a call to
+	    <c>erlang:dist_ctrl_get_data(<anno>DHandle</anno>)</c>
+	    will return <c>Data</c> to pass over the channel as well as
+	    the <c>Size</c> of <c>Data</c> in bytes. This is returned as
+	    a tuple on the form <c>{Size, Data}</c>.
+	  </item>
+	</taglist>
+	<p>All options are set to default when a channel is closed.</p>
+	<note><p>
+	  Only the process registered as distribution
+	  controller for the distribution channel identified by
+	  <c><anno>DHandle</anno></c> is allowed to call this
+	  function.
+	</p></note>
+	<p>
+	  This function is used when implementing an alternative
+	  distribution carrier using processes as distribution
+	  controllers. <c><anno>DHandle</anno></c> is retrived
+	  via the callback
+	  <seealso marker="erts:alt_dist#hs_data_f_handshake_complete"><c>f_handshake_complete</c></seealso>.
+	  More information can be found in the documentation of
+	  <seealso marker="erts:alt_dist#distribution_module">ERTS
+	  User's Guide ➜ How to implement an Alternative Carrier
+	  for the Erlang Distribution ➜ Distribution Module</seealso>.
+	</p>
+      </desc>
+    </func>
+    
+    <func>
       <name name="element" arity="2" since=""/>
       <fsummary>Return the Nth element of a tuple.</fsummary>
       <type_desc variable="N">1..tuple_size(<anno>Tuple</anno>)</type_desc>
diff --git a/erts/emulator/beam/atom.names b/erts/emulator/beam/atom.names
index f81082a698..412d689246 100644
--- a/erts/emulator/beam/atom.names
+++ b/erts/emulator/beam/atom.names
@@ -290,6 +290,7 @@ atom Ge='>='
 atom generational
 atom get_all_trap
 atom get_seq_token
+atom get_size
 atom get_tcw
 atom gather_gc_info_result
 atom gather_io_bytes
diff --git a/erts/emulator/beam/bif.tab b/erts/emulator/beam/bif.tab
index 11941db8cd..34a0be4f2d 100644
--- a/erts/emulator/beam/bif.tab
+++ b/erts/emulator/beam/bif.tab
@@ -160,6 +160,8 @@ bif erlang:dist_ctrl_input_handler/2
 bif erlang:dist_ctrl_put_data/2
 bif erlang:dist_ctrl_get_data/1
 bif erlang:dist_ctrl_get_data_notification/1
+bif erlang:dist_ctrl_get_opt/2
+bif erlang:dist_ctrl_set_opt/3
 
 # Static native functions in erts_internal
 bif erts_internal:port_info/1
diff --git a/erts/emulator/beam/dist.c b/erts/emulator/beam/dist.c
index b50c8273b1..ff5f766de7 100644
--- a/erts/emulator/beam/dist.c
+++ b/erts/emulator/beam/dist.c
@@ -3253,6 +3253,86 @@ dist_ctrl_put_data_2(BIF_ALIST_2)
 }
 
 BIF_RETTYPE
+dist_ctrl_set_opt_3(BIF_ALIST_3)
+{
+    DistEntry *dep = ERTS_PROC_GET_DIST_ENTRY(BIF_P);
+    Uint32 conn_id;
+    BIF_RETTYPE ret;
+
+    if (!dep)
+        BIF_ERROR(BIF_P, EXC_NOTSUP);
+
+    if (erts_dhandle_to_dist_entry(BIF_ARG_1, &conn_id) != dep)
+        BIF_ERROR(BIF_P, BADARG);
+
+    erts_de_rlock(dep);
+
+    if (dep->connection_id != conn_id)
+        ERTS_BIF_PREP_ERROR(ret, BIF_P, BADARG);
+    else {
+
+        switch (BIF_ARG_2) {
+        case am_get_size:
+            ERTS_BIF_PREP_RET(ret, (dep->opts & ERTS_DIST_CTRL_OPT_GET_SIZE
+                                    ? am_true
+                                    : am_false));
+            if (BIF_ARG_3 == am_true) 
+                dep->opts |= ERTS_DIST_CTRL_OPT_GET_SIZE;
+            else if (BIF_ARG_3 == am_false)
+                dep->opts &= ~ERTS_DIST_CTRL_OPT_GET_SIZE;
+            else
+                ERTS_BIF_PREP_ERROR(ret, BIF_P, BADARG);
+            break;
+        default:
+            ERTS_BIF_PREP_ERROR(ret, BIF_P, BADARG);
+            break;
+        }
+
+    }
+
+    erts_de_runlock(dep);
+
+    return ret;
+}
+
+BIF_RETTYPE
+dist_ctrl_get_opt_2(BIF_ALIST_2)
+{
+    DistEntry *dep = ERTS_PROC_GET_DIST_ENTRY(BIF_P);
+    Uint32 conn_id;
+    BIF_RETTYPE ret;
+
+    if (!dep)
+        BIF_ERROR(BIF_P, EXC_NOTSUP);
+
+    if (erts_dhandle_to_dist_entry(BIF_ARG_1, &conn_id) != dep)
+        BIF_ERROR(BIF_P, BADARG);
+
+    erts_de_rlock(dep);
+
+    if (dep->connection_id != conn_id)
+        ERTS_BIF_PREP_ERROR(ret, BIF_P, BADARG);
+    else {
+
+        switch (BIF_ARG_2) {
+        case am_get_size:
+            ERTS_BIF_PREP_RET(ret, (dep->opts & ERTS_DIST_CTRL_OPT_GET_SIZE
+                                    ? am_true
+                                    : am_false));
+            break;
+        default:
+            ERTS_BIF_PREP_ERROR(ret, BIF_P, BADARG);
+            break;
+        }
+
+    }
+
+    erts_de_runlock(dep);
+
+    return ret;
+}
+    
+BIF_RETTYPE
 dist_get_stat_1(BIF_ALIST_1)
 {
     Sint64 read, write, pend;
@@ -3332,7 +3412,9 @@ dist_ctrl_get_data_1(BIF_ALIST_1)
     Eterm *hp;
     ProcBin *pb;
     erts_aint_t qsize;
-    Uint32 conn_id;
+    Uint32 conn_id, get_size;
+    Eterm res;
+    Uint hsz, bin_sz;
 
     if (!dep)
         BIF_ERROR(BIF_P, EXC_NOTSUP);
@@ -3400,15 +3482,26 @@ dist_ctrl_get_data_1(BIF_ALIST_1)
 
     erts_de_runlock(dep);
 
-    hp =  HAlloc(BIF_P, PROC_BIN_SIZE);
+    bin_sz = obuf->ext_endp - obuf->extp;
+    hsz = PROC_BIN_SIZE;
+
+    get_size = dep->opts & ERTS_DIST_CTRL_OPT_GET_SIZE;
+    if (get_size) {
+        hsz += 3; /* 2 tuple */
+        if (!IS_USMALL(0, bin_sz))
+            hsz += BIG_UINT_HEAP_SIZE;
+    }
+
+    hp = HAlloc(BIF_P, hsz);
     pb = (ProcBin *) (char *) hp;
     pb->thing_word = HEADER_PROC_BIN;
-    pb->size = obuf->ext_endp - obuf->extp;
+    pb->size = bin_sz;
     pb->next = MSO(BIF_P).first;
     MSO(BIF_P).first = (struct erl_off_heap_header*) pb;
     pb->val = ErtsDistOutputBuf2Binary(obuf);
     pb->bytes = (byte*) obuf->extp;
     pb->flags = 0;
+    hp += PROC_BIN_SIZE;
 
     qsize = erts_atomic_add_read_nob(&dep->qsize, -size_obuf(obuf));
     ASSERT(qsize >= 0);
@@ -3425,7 +3518,20 @@ dist_ctrl_get_data_1(BIF_ALIST_1)
         }
     }
 
-    BIF_RET2(make_binary(pb), (initial_reds - reds));
+    res = make_binary(pb);
+
+    if (get_size) {
+        Eterm sz_term;
+        if (IS_USMALL(0, bin_sz))
+            sz_term = make_small(bin_sz);
+        else {
+            sz_term = uint_to_big(bin_sz, hp);
+            hp += BIG_UINT_HEAP_SIZE;
+        }
+        res = TUPLE2(hp, sz_term, res);
+    }
+
+    BIF_RET2(res, (initial_reds - reds));
 }
 
 void
diff --git a/erts/emulator/beam/dist.h b/erts/emulator/beam/dist.h
index c4bb967592..5bd22cc31f 100644
--- a/erts/emulator/beam/dist.h
+++ b/erts/emulator/beam/dist.h
@@ -172,6 +172,9 @@ extern int erts_is_alive;
 /* Pending connection; signals can be enqueued */
 #define ERTS_DSIG_PREP_PENDING	        4
 
+/* dist_ctrl_{g,s}et_option/2 */
+#define ERTS_DIST_CTRL_OPT_GET_SIZE     ((Uint32) (1 << 0))
+
 #ifdef DEBUG
 #define ERTS_DBG_CHK_NO_DIST_LNK(D, R, L) \
     erts_dbg_chk_no_dist_proc_link((D), (R), (L))
diff --git a/erts/emulator/beam/erl_node_tables.c b/erts/emulator/beam/erl_node_tables.c
index afafaf48dc..1adb101e30 100644
--- a/erts/emulator/beam/erl_node_tables.c
+++ b/erts/emulator/beam/erl_node_tables.c
@@ -177,6 +177,7 @@ dist_table_alloc(void *dep_tmpl)
     dep->connection_id			= 0;
     dep->state				= ERTS_DE_STATE_IDLE;
     dep->flags				= 0;
+    dep->opts                           = 0;
     dep->version			= 0;
 
     dep->mld                            = NULL;
@@ -659,6 +660,7 @@ erts_set_dist_entry_not_connected(DistEntry *dep)
 
     dep->state = ERTS_DE_STATE_IDLE;
     dep->flags = 0;
+    dep->opts = 0;
     dep->prev = NULL;
     dep->cid = NIL;
 
diff --git a/erts/emulator/beam/erl_node_tables.h b/erts/emulator/beam/erl_node_tables.h
index d5daf0c2df..8153699596 100644
--- a/erts/emulator/beam/erl_node_tables.h
+++ b/erts/emulator/beam/erl_node_tables.h
@@ -148,6 +148,7 @@ struct dist_entry_ {
     enum dist_entry_state state;
     Uint32 flags;		/* Distribution flags, like hidden, 
 				   atom cache etc. */
+    Uint32 opts;
     unsigned long version;	/* Protocol version */
 
     ErtsMonLnkDist *mld;        /* Monitors and links */
diff --git a/erts/preloaded/ebin/erlang.beam b/erts/preloaded/ebin/erlang.beam
index 1b0cb5b50c..62dc8702e7 100644
--- a/erts/preloaded/ebin/erlang.beam
+++ b/erts/preloaded/ebin/erlang.beam
diff --git a/erts/preloaded/src/erlang.erl b/erts/preloaded/src/erlang.erl
index a5b60cc845..ac73946dc0 100644
--- a/erts/preloaded/src/erlang.erl
+++ b/erts/preloaded/src/erlang.erl
@@ -48,6 +48,8 @@
          dist_ctrl_put_data/2,
          dist_ctrl_get_data/1,
          dist_ctrl_get_data_notification/1,
+         dist_ctrl_get_opt/2,
+         dist_ctrl_set_opt/3,
          dist_get_stat/1]).
 
 -deprecated([get_stacktrace/0,now/0]).
@@ -3326,7 +3328,8 @@ dist_ctrl_input_handler(_DHandle, _InputHandler) ->
 dist_ctrl_put_data(_DHandle, _Data) ->
     erlang:nif_error(undefined).
 
--spec erlang:dist_ctrl_get_data(DHandle) -> Data | 'none' when
+-spec erlang:dist_ctrl_get_data(DHandle) -> {Size, Data} | Data | 'none' when
+      Size :: non_neg_integer(),
       DHandle :: dist_handle(),
       Data :: iodata().
 
@@ -3339,6 +3342,21 @@ dist_ctrl_get_data(_DHandle) ->
 dist_ctrl_get_data_notification(_DHandle) ->
     erlang:nif_error(undefined).
 
+-spec erlang:dist_ctrl_set_opt(DHandle, 'get_size', Value) -> OldValue when
+      DHandle :: dist_handle(),
+      Value :: boolean(),
+      OldValue :: boolean().
+
+dist_ctrl_set_opt(_DHandle, _Opt, _Val) ->
+    erlang:nif_error(undefined).
+
+-spec erlang:dist_ctrl_get_opt(DHandle, 'get_size') -> Value when
+      DHandle :: dist_handle(),
+      Value :: boolean().
+
+dist_ctrl_get_opt(_DHandle, _Opt) ->
+    erlang:nif_error(undefined).
+
 -spec erlang:dist_get_stat(DHandle) -> Res when
       DHandle :: dist_handle(),
       InputPackets :: non_neg_integer(),
diff --git a/lib/tools/emacs/erlang.el b/lib/tools/emacs/erlang.el
index 3cbe9daa60..38c0eba92b 100644
--- a/lib/tools/emacs/erlang.el
+++ b/lib/tools/emacs/erlang.el
@@ -905,8 +905,10 @@ resulting regexp is surrounded by \\_< and \\_>."
       "dist_get_stat"
       "dist_ctrl_get_data"
       "dist_ctrl_get_data_notification"
+      "dist_ctrl_get_opt"
       "dist_ctrl_input_handler"
       "dist_ctrl_put_data"
+      "dist_ctrl_set_opt"
       "dmonitor_node"
       "dt_append_vm_tag_data"
       "dt_get_tag"