content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.server.util;
  
  import org.apache.kafka.clients.ClientRequest;
  import org.apache.kafka.clients.ClientResponse;
  import org.apache.kafka.clients.KafkaClient;
  import org.apache.kafka.clients.RequestCompletionHandler;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.errors.AuthenticationException;
  import org.apache.kafka.common.errors.DisconnectException;
  import org.apache.kafka.common.internals.FatalExitError;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.requests.AbstractRequest;
  import org.apache.kafka.common.utils.Time;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.ValueSource;
  import org.mockito.ArgumentMatchers;
  
  import java.io.IOException;
  import java.util.ArrayDeque;
  import java.util.Collection;
  import java.util.List;
  import java.util.Queue;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Consumer;
  
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertInstanceOf;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.anyLong;
  import static org.mockito.ArgumentMatchers.same;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.verifyNoMoreInteractions;
  import static org.mockito.Mockito.when;
  
  public class InterBrokerSendThreadTest {
  
      private final Time time = new MockTime();
      private final KafkaClient networkClient = mock(KafkaClient.class);
      private final StubCompletionHandler completionHandler = new StubCompletionHandler();
      private final int requestTimeoutMs = 1000;
  
      class TestInterBrokerSendThread extends InterBrokerSendThread {
  
          private final Consumer<Throwable> exceptionCallback;
          private final Queue<RequestAndCompletionHandler> queue = new ArrayDeque<>();
  
          TestInterBrokerSendThread() {
              this(
                  InterBrokerSendThreadTest.this.networkClient,
                  t -> {
                      throw (t instanceof RuntimeException)
                          ? ((RuntimeException) t)
                          : new RuntimeException(t);
                  });
          }
  
          TestInterBrokerSendThread(KafkaClient networkClient, Consumer<Throwable> exceptionCallback) {
              super("name", networkClient, requestTimeoutMs, time);
              this.exceptionCallback = exceptionCallback;
          }
  
          void enqueue(RequestAndCompletionHandler request) {
              queue.offer(request);
          }
  
          @Override
          public Collection<RequestAndCompletionHandler> generateRequests() {
              return queue.isEmpty() ? List.of() : List.of(queue.poll());
          }
  
          @Override
          protected void pollOnce(long maxTimeoutMs) {
              try {
                  super.pollOnce(maxTimeoutMs);
              } catch (Throwable t) {
                  exceptionCallback.accept(t);
              }
          }
      }
  
      @Test
      public void testShutdownThreadShouldNotCauseException() throws InterruptedException, IOException {
          // InterBrokerSendThread#shutdown calls NetworkClient#initiateClose first so NetworkClient#poll
          // can throw DisconnectException when thread is running
          when(networkClient.poll(anyLong(), anyLong())).thenThrow(new DisconnectException());
          when(networkClient.active()).thenReturn(false);
  
          AtomicReference<Throwable> exception = new AtomicReference<>();
          final InterBrokerSendThread thread =
              new TestInterBrokerSendThread(networkClient, exception::getAndSet);
          thread.shutdown();
          thread.pollOnce(100);
  
          verify(networkClient).poll(anyLong(), anyLong());
          verify(networkClient).initiateClose();
          verify(networkClient).close();
          verify(networkClient).active();
          verifyNoMoreInteractions(networkClient);
  
          assertNull(exception.get());
      }
  
      @Test
      public void testDisconnectWithoutShutdownShouldCauseException() {
          DisconnectException de = new DisconnectException();
          when(networkClient.poll(anyLong(), anyLong())).thenThrow(de);
          when(networkClient.active()).thenReturn(true);
  
          AtomicReference<Throwable> throwable = new AtomicReference<>();
          final InterBrokerSendThread thread =
              new TestInterBrokerSendThread(networkClient, throwable::getAndSet);
          thread.pollOnce(100);
  
          verify(networkClient).poll(anyLong(), anyLong());
          verify(networkClient).active();
          verifyNoMoreInteractions(networkClient);
  
          Throwable thrown = throwable.get();
          assertNotNull(thrown);
          assertInstanceOf(FatalExitError.class, thrown);
      }
  
      @Test
      public void testShouldNotSendAnythingWhenNoRequests() {
          final InterBrokerSendThread sendThread = new TestInterBrokerSendThread();
  
          // poll is always called but there should be no further invocations on NetworkClient
          when(networkClient.poll(anyLong(), anyLong())).thenReturn(List.of());
  
          sendThread.doWork();
  
          verify(networkClient).poll(anyLong(), anyLong());
          verifyNoMoreInteractions(networkClient);
  
          assertFalse(completionHandler.executedWithDisconnectedResponse);
      }
  
      @Test
      public void testShouldCreateClientRequestAndSendWhenNodeIsReady() {
          final AbstractRequest.Builder<?> request = new StubRequestBuilder<>();
          final Node node = new Node(1, "", 8080);
          final RequestAndCompletionHandler handler =
              new RequestAndCompletionHandler(time.milliseconds(), node, request, completionHandler);
          final TestInterBrokerSendThread sendThread = new TestInterBrokerSendThread();
  
          final ClientRequest clientRequest =
              new ClientRequest("dest", request, 0, "1", 0, true, requestTimeoutMs, handler.handler);
  
          when(networkClient.newClientRequest(
              ArgumentMatchers.eq("1"),
              same(handler.request),
              anyLong(),
              ArgumentMatchers.eq(true),
              ArgumentMatchers.eq(requestTimeoutMs),
              same(handler.handler)
          )).thenReturn(clientRequest);
  
          when(networkClient.ready(node, time.milliseconds())).thenReturn(true);
  
          when(networkClient.poll(anyLong(), anyLong())).thenReturn(List.of());
  
          sendThread.enqueue(handler);
          sendThread.doWork();
  
          verify(networkClient)
              .newClientRequest(
                  ArgumentMatchers.eq("1"),
                  same(handler.request),
                  anyLong(),
                  ArgumentMatchers.eq(true),
                  ArgumentMatchers.eq(requestTimeoutMs),
                  same(handler.handler));
          verify(networkClient).ready(any(), anyLong());
          verify(networkClient).send(same(clientRequest), anyLong());
          verify(networkClient).poll(anyLong(), anyLong());
          verifyNoMoreInteractions(networkClient);
  
          assertFalse(completionHandler.executedWithDisconnectedResponse);
      }
  
      @Test
      public void testShouldCallCompletionHandlerWithDisconnectedResponseWhenNodeNotReady() {
          final AbstractRequest.Builder<?> request = new StubRequestBuilder<>();
          final Node node = new Node(1, "", 8080);
          final RequestAndCompletionHandler handler =
              new RequestAndCompletionHandler(time.milliseconds(), node, request, completionHandler);
          final TestInterBrokerSendThread sendThread = new TestInterBrokerSendThread();
  
          final ClientRequest clientRequest =
              new ClientRequest("dest", request, 0, "1", 0, true, requestTimeoutMs, handler.handler);
  
          when(networkClient.newClientRequest(
              ArgumentMatchers.eq("1"),
              same(handler.request),
              anyLong(),
              ArgumentMatchers.eq(true),
              ArgumentMatchers.eq(requestTimeoutMs),
              same(handler.handler)
          )).thenReturn(clientRequest);
  
          when(networkClient.ready(node, time.milliseconds())).thenReturn(false);
  
          when(networkClient.connectionDelay(any(), anyLong())).thenReturn(0L);
  
          when(networkClient.poll(anyLong(), anyLong())).thenReturn(List.of());
  
          when(networkClient.connectionFailed(node)).thenReturn(true);
  
          when(networkClient.authenticationException(node)).thenReturn(new AuthenticationException(""));
  
          sendThread.enqueue(handler);
          sendThread.doWork();
  
          verify(networkClient)
              .newClientRequest(
                  ArgumentMatchers.eq("1"),
                  same(handler.request),
                  anyLong(),
                  ArgumentMatchers.eq(true),
                  ArgumentMatchers.eq(requestTimeoutMs),
                  same(handler.handler));
          verify(networkClient).ready(any(), anyLong());
          verify(networkClient).connectionDelay(any(), anyLong());
          verify(networkClient).poll(anyLong(), anyLong());
          verify(networkClient).connectionFailed(any());
          verify(networkClient).authenticationException(any());
          verifyNoMoreInteractions(networkClient);
  
          assertTrue(completionHandler.executedWithDisconnectedResponse);
      }
  
      @Test
      public void testFailingExpiredRequests() {
          final AbstractRequest.Builder<?> request = new StubRequestBuilder<>();
          final Node node = new Node(1, "", 8080);
          final RequestAndCompletionHandler handler =
              new RequestAndCompletionHandler(time.milliseconds(), node, request, completionHandler);
          final TestInterBrokerSendThread sendThread = new TestInterBrokerSendThread();
  
          final ClientRequest clientRequest =
              new ClientRequest(
                  "dest", request, 0, "1", time.milliseconds(), true, requestTimeoutMs, handler.handler);
          time.sleep(1500L);
  
          when(networkClient.newClientRequest(
              ArgumentMatchers.eq("1"),
              same(handler.request),
              ArgumentMatchers.eq(handler.creationTimeMs),
              ArgumentMatchers.eq(true),
              ArgumentMatchers.eq(requestTimeoutMs),
              same(handler.handler)
          )).thenReturn(clientRequest);
  
          // make the node unready so the request is not cleared
          when(networkClient.ready(node, time.milliseconds())).thenReturn(false);
  
          when(networkClient.connectionDelay(any(), anyLong())).thenReturn(0L);
  
          when(networkClient.poll(anyLong(), anyLong())).thenReturn(List.of());
  
          // rule out disconnects so the request stays for the expiry check
          when(networkClient.connectionFailed(node)).thenReturn(false);
  
          sendThread.enqueue(handler);
          sendThread.doWork();
  
          verify(networkClient)
              .newClientRequest(
                  ArgumentMatchers.eq("1"),
                  same(handler.request),
                  ArgumentMatchers.eq(handler.creationTimeMs),
                  ArgumentMatchers.eq(true),
                  ArgumentMatchers.eq(requestTimeoutMs),
                  same(handler.handler));
          verify(networkClient).ready(any(), anyLong());
          verify(networkClient).connectionDelay(any(), anyLong());
          verify(networkClient).poll(anyLong(), anyLong());
          verify(networkClient).connectionFailed(any());
          verifyNoMoreInteractions(networkClient);
  
          assertFalse(sendThread.hasUnsentRequests());
          assertTrue(completionHandler.executedWithDisconnectedResponse);
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testInterruption(boolean isShuttingDown) throws InterruptedException, IOException {
          Exception interrupted = new InterruptedException();
  
          // InterBrokerSendThread#shutdown calls NetworkClient#initiateClose first so NetworkClient#poll
          // can throw InterruptedException if a callback request that throws it is handled
          when(networkClient.poll(anyLong(), anyLong())).thenAnswer(t -> {
              throw interrupted;
          });
  
          AtomicReference<Throwable> exception = new AtomicReference<>();
          final InterBrokerSendThread thread =
              new TestInterBrokerSendThread(networkClient, t -> {
                  if (isShuttingDown)
                      assertInstanceOf(InterruptedException.class, t);
                  else
                      assertInstanceOf(FatalExitError.class, t);
                  exception.getAndSet(t);
              });
  
          if (isShuttingDown)
              thread.shutdown();
          thread.pollOnce(100);
  
          verify(networkClient).poll(anyLong(), anyLong());
          if (isShuttingDown) {
              verify(networkClient).initiateClose();
              verify(networkClient).close();
          }
          verifyNoMoreInteractions(networkClient);
          assertNotNull(exception.get());
      }
  
      private static class StubRequestBuilder<T extends AbstractRequest>
          extends AbstractRequest.Builder<T> {
  
          private StubRequestBuilder() {
              super(ApiKeys.END_TXN);
          }
  
          @Override
          public T build(short version) {
              return null;
          }
      }
  
      private static class StubCompletionHandler implements RequestCompletionHandler {
  
          public boolean executedWithDisconnectedResponse = false;
          ClientResponse response = null;
  
          @Override
          public void onComplete(ClientResponse response) {
              this.executedWithDisconnectedResponse = response.wasDisconnected();
              this.response = response;
          }
      }
  }
  