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

6 files changed, 119 insertions, 4 deletions

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 */