Merge branch 'rickard/literal-whole-message/OTP-17307' into maint-21

* rickard/literal-whole-message/OTP-17307:
  Fix literal collection of messages

2 files changed, 74 insertions, 6 deletions

diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index 740b4ca022..1c5ed9d30b 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -902,7 +902,7 @@ msg_copy_literal_area(ErtsMessage *msgp, int *redsp,
 
     *redsp += 1;
 
-    if (!ERTS_SIG_IS_INTERNAL_MSG(msgp) || !msgp->data.attached)
+    if (!ERTS_SIG_IS_INTERNAL_MSG(msgp))
         return;
 
     if (msgp->data.attached == ERTS_MSG_COMBINED_HFRAG)
@@ -910,6 +910,23 @@ msg_copy_literal_area(ErtsMessage *msgp, int *redsp,
     else
         hfrag = msgp->data.heap_frag;
 
+    /*
+     * Literals should only be able to appear in the
+     * first message reference, i.e., the message
+     * itself...
+     */
+    if (ErtsInArea(msgp->m[0], literals, lit_bsize))
+        lit_sz += size_object(msgp->m[0]);
+
+#ifdef DEBUG
+    {
+        int i;
+        for (i = 1; i < ERL_MESSAGE_REF_ARRAY_SZ; i++) {
+            ASSERT(!ErtsInArea(msgp->m[i], literals, lit_bsize));
+        }
+    }
+#endif
+
     for (hf = hfrag; hf; hf = hf->next) {
         lit_sz += hfrag_literal_size(&hf->mem[0],
                                      &hf->mem[hf->used_size],
@@ -922,6 +939,11 @@ msg_copy_literal_area(ErtsMessage *msgp, int *redsp,
         ErlHeapFragment *bp = new_message_buffer(lit_sz);
         Eterm *hp = bp->mem;
 
+        if (ErtsInArea(msgp->m[0], literals, lit_bsize)) {
+            Uint sz = size_object(msgp->m[0]);
+            msgp->m[0] = copy_struct(msgp->m[0], sz, &hp, &bp->off_heap);
+        }
+
         for (hf = hfrag; hf; hf = hf->next) {
             hfrag_literal_copy(&hp, &bp->off_heap,
                                &hf->mem[0],
@@ -930,10 +952,14 @@ msg_copy_literal_area(ErtsMessage *msgp, int *redsp,
             hfrag = hf;
         }
 
-        /* link new hfrag last */
-        ASSERT(hfrag->next == NULL);
-        hfrag->next = bp;
         bp->next = NULL;
+        /* link new hfrag last */
+        if (!hfrag)
+            msgp->data.heap_frag = bp;
+        else {
+            ASSERT(hfrag->next == NULL);
+            hfrag->next = bp;
+        }
     }
 }
 
diff --git a/erts/emulator/test/persistent_term_SUITE.erl b/erts/emulator/test/persistent_term_SUITE.erl
index 93eb026ced..7c0b1ab3db 100644
--- a/erts/emulator/test/persistent_term_SUITE.erl
+++ b/erts/emulator/test/persistent_term_SUITE.erl
@@ -25,7 +25,7 @@
 	 basic/1,purging/1,sharing/1,get_trapping/1,
          info/1,info_trapping/1,killed_while_trapping/1,
          off_heap_values/1,keys/1,collisions/1,
-         init_restart/1]).
+         init_restart/1,whole_message/1]).
 
 %%
 -export([test_init_restart_cmd/1]).
@@ -37,7 +37,7 @@ suite() ->
 all() ->
     [basic,purging,sharing,get_trapping,info,info_trapping,
      killed_while_trapping,off_heap_values,keys,collisions,
-     init_restart].
+     init_restart,whole_message].
 
 init_per_suite(Config) ->
     %% Put a term in the dict so that we know that the testcases handle
@@ -596,6 +596,48 @@ do_test_init_restart_cmd(File) ->
             init:stop()
     end.
 
+%% Test that the literal is copied when removed also when
+%% the whole message is a literal...
+
+whole_message(Config) when is_list(Config) ->
+    whole_message_test(on_heap),
+    whole_message_test(off_heap),
+    ok.
+    
+whole_message_test(MQD) ->
+    io:format("Testing on ~p~n", [MQD]),
+    Go = make_ref(),
+    Done = make_ref(),
+    TestRef = make_ref(),
+    Tester = self(),
+    persistent_term:put(test_ref, TestRef),
+    Pid = spawn_opt(fun () ->
+                             receive Go -> ok end,
+                             receive TestRef -> ok end,
+                             receive TestRef -> ok end,
+                             receive TestRef -> ok end,
+                             receive [TestRef] -> ok end,
+                             receive [TestRef] -> ok end,
+                             receive [TestRef] -> ok end,
+                             Tester ! Done
+                     end, [link, {message_queue_data, MQD}]),
+    Pid ! persistent_term:get(test_ref),
+    Pid ! persistent_term:get(test_ref),
+    Pid ! persistent_term:get(test_ref),
+    %% Throw in some messages with a reference from the heap
+    %% while we're at it...
+    Pid ! [persistent_term:get(test_ref)],
+    Pid ! [persistent_term:get(test_ref)],
+    Pid ! [persistent_term:get(test_ref)],
+    persistent_term:erase(test_ref),
+    receive after 1000 -> ok end,
+    Pid ! Go,
+    receive Done -> ok end,
+    unlink(Pid),
+    exit(Pid, kill),
+    false = is_process_alive(Pid),
+    ok.
+
 %% Check that there is the same number of persistents terms before
 %% and after each test case.