Repository: apache/flink
File: flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/mailbox/TaskMailboxImplTest.java (Java)

Code before:
                                                    taskMailbox.put(new Mail(() -> {}, 1, "mailD"));
                                                    Thread.sleep(1);
                                                } catch (Exception e) {
                                                }
                                            }
                                        }))
                .start();

        exclusiveCodeStarted.await();
        // make sure that all 10 messages have been actually enqueued.
        assertThat(taskMailbox.close()).hasSize(numMails);
    }

    /**
     * We expect the high priority mail can be taken immediately during taking mail from batch
     * queue.
     */
    @Test
    void testPutHighPriorityMailDuringTakingFromBatch() throws Exception {
        Mail mailA = new Mail(() -> {}, DEFAULT_PRIORITY, "mailA");
        Mail mailB = new Mail(() -> {}, DEFAULT_PRIORITY, "mailB");
        Mail mailC =
                new Mail(
                        MailboxExecutor.MailOptions.highPriority(),
                        () -> {},
                        DEFAULT_PRIORITY,
                        "mailC");

        Thread thread = new Thread(() -> taskMailbox.runExclusively(() -> taskMailbox.put(mailC)));

        taskMailbox.put(mailA);
        taskMailbox.put(mailB);
        assertThat(taskMailbox.createBatch()).isTrue();
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailA);

        thread.start();
        thread.join();

        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailC);
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailB);
    }
}


Review comment (suggestion):
!nit, this is confusing as this message is high priority, I'd rather use Interger.MAX_INT with comment that this is ignored

Code after:

        Mail mail1 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail1");
        Mail mail2 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail2");
        Mail mail3 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail3");
        Mail mail4 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail4");
        Mail mail5 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail5");
        Mail mail6 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail6");
        Mail mail7 = new Mail(() -> {}, DEFAULT_PRIORITY, "mail7");

        putMail(mail1, isAsyncPut);
        putMail(mail2, isAsyncPut);
        putMail(mail3, isAsyncPut);
        putMail(mail4, isAsyncPut);
        assertThat(taskMailbox.createBatch()).isTrue();
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mail1);

        putMail(mailA, isAsyncPut);
        putMail(mailB, isAsyncPut);
        putMail(mailC, isAsyncPut);
        // All urgent mails should be taken first, including all types of take methods
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailA);
        assertThat(taskMailbox.tryTake(DEFAULT_PRIORITY)).hasValue(mailB);
        assertThat(taskMailbox.take(DEFAULT_PRIORITY)).isEqualTo(mailC);

        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mail2);

        putMail(mail5, isAsyncPut);
        putMail(mail6, isAsyncPut);
        putMail(mail7, isAsyncPut);
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mail3);

        putMail(mailD, isAsyncPut);
        putMail(mailE, isAsyncPut);
        assertThat(taskMailbox.createBatch()).isTrue();
        putMail(mailF, isAsyncPut);
        putMail(mailG, isAsyncPut);

        // All urgent mails should be taken first, including all types of take methods
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailD);
        assertThat(taskMailbox.tryTake(DEFAULT_PRIORITY)).hasValue(mailE);
        assertThat(taskMailbox.take(DEFAULT_PRIORITY)).isEqualTo(mailF);
        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mailG);

        assertThat(taskMailbox.tryTakeFromBatch()).hasValue(mail4);
        assertThat(taskMailbox.tryTake(DEFAULT_PRIORITY)).hasValue(mail5);
        assertThat(taskMailbox.take(DEFAULT_PRIORITY)).isEqualTo(mail6);
        assertThat(taskMailbox.tryTake(DEFAULT_PRIORITY)).hasValue(mail7);
    }

    @ValueSource(booleans = {true, false})
    @ParameterizedTest