1 files changed, 21 insertions, 13 deletions

diff --git a/erts/emulator/beam/io.c b/erts/emulator/beam/io.c
index 133ab485d9..5325480901 100644
--- a/erts/emulator/beam/io.c
+++ b/erts/emulator/beam/io.c
@@ -375,7 +375,6 @@ static Port *create_port(char *name,
     prt->control_flags = 0;
     prt->bytes_in = 0;
     prt->bytes_out = 0;
-    prt->dist_entry = NULL;
     ERTS_PORT_INIT_CONNECTED(prt, pid);
     prt->common.u.alive.reg = NULL;
     ERTS_PTMR_INIT(prt);
@@ -3613,12 +3612,12 @@ terminate_port(Port *prt)
 
     erts_cleanup_port_data(prt);
 
+    ASSERT(erts_prtsd_get(prt, ERTS_PRTSD_DIST_ENTRY) == NULL);
+
     psd = (ErtsPrtSD *) erts_atomic_read_nob(&prt->psd);
     if (psd)
 	erts_free(ERTS_ALC_T_PRTSD, psd);
 
-    ASSERT(prt->dist_entry == NULL);
-
     kill_port(prt);
 
     /*
@@ -3759,10 +3758,12 @@ erts_deliver_port_exit(Port *prt, Eterm from, Eterm reason, int send_closed,
        DRV_MONITOR_UNLOCK_PDL(prt);
    }
 
-   if ((state & ERTS_PORT_SFLG_DISTRIBUTION) && prt->dist_entry) {
-       erts_do_net_exits(prt->dist_entry, modified_reason);
-       erts_deref_dist_entry(prt->dist_entry);
-       prt->dist_entry = NULL;
+   if (state & ERTS_PORT_SFLG_DISTRIBUTION) {
+       DistEntry *dep = (DistEntry*) erts_prtsd_get(prt, ERTS_PRTSD_DIST_ENTRY);
+       ASSERT(dep);
+       erts_do_net_exits(dep, modified_reason);
+       erts_deref_dist_entry(dep);
+       erts_prtsd_set(prt, ERTS_PRTSD_DIST_ENTRY, NULL);
        erts_atomic32_read_band_relb(&prt->state,
 				    ~ERTS_PORT_SFLG_DISTRIBUTION);
    }
@@ -5050,7 +5051,7 @@ set_busy_port(ErlDrvPort dprt, int on)
             DTRACE1(port_not_busy, port_str);
         }
 #endif
-	if (prt->dist_entry) {
+	if (erts_prtsd_get(prt, ERTS_PRTSD_DIST_ENTRY) != NULL) {
 	    /*
 	     * Processes suspended on distribution ports are
 	     * normally queued on the dist entry.
@@ -6169,9 +6170,12 @@ int driver_output_binary(ErlDrvPort ix, char* hbuf, ErlDrvSizeT hlen,
     else
 	erts_atomic64_add_nob(&bytes_in, (erts_aint64_t) (hlen + len));
     if (state & ERTS_PORT_SFLG_DISTRIBUTION) {
-        erts_atomic64_inc_nob(&prt->dist_entry->in);
+        DistEntry* dep = (DistEntry*) erts_prtsd_get(prt, ERTS_PRTSD_DIST_ENTRY);
+        Uint32 conn_id = (Uint32)(UWord) erts_prtsd_get(prt, ERTS_PRTSD_CONN_ID);
+        erts_atomic64_inc_nob(&dep->in);
 	return erts_net_message(prt,
-				prt->dist_entry,
+				dep,
+                                conn_id,
 				(byte*) hbuf, hlen,
 				(byte*) (bin->orig_bytes+offs), len);
     }
@@ -6210,15 +6214,19 @@ int driver_output2(ErlDrvPort ix, char* hbuf, ErlDrvSizeT hlen,
     else
 	erts_atomic64_add_nob(&bytes_in, (erts_aint64_t) (hlen + len));
     if (state & ERTS_PORT_SFLG_DISTRIBUTION) {
-        erts_atomic64_inc_nob(&prt->dist_entry->in);
+        DistEntry *dep = (DistEntry*) erts_prtsd_get(prt, ERTS_PRTSD_DIST_ENTRY);
+        Uint32 conn_id = (Uint32)(UWord) erts_prtsd_get(prt, ERTS_PRTSD_CONN_ID);
+        erts_atomic64_inc_nob(&dep->in);
 	if (len == 0)
 	    return erts_net_message(prt,
-				    prt->dist_entry,
+				    dep,
+                                    conn_id,
 				    NULL, 0,
 				    (byte*) hbuf, hlen);
 	else
 	    return erts_net_message(prt,
-				    prt->dist_entry,
+				    dep,
+                                    conn_id,
 				    (byte*) hbuf, hlen,
 				    (byte*) buf, len);
     }