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.connect.storage;
  
  import org.apache.kafka.connect.util.Callback;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.mockito.ArgumentCaptor;
  import org.mockito.junit.jupiter.MockitoExtension;
  import org.mockito.junit.jupiter.MockitoSettings;
  import org.mockito.quality.Strictness;
  
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.mockito.ArgumentMatchers.eq;
  import static org.mockito.ArgumentMatchers.isNull;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.verifyNoInteractions;
  import static org.mockito.Mockito.when;
  
  @ExtendWith(MockitoExtension.class)
  @MockitoSettings(strictness = Strictness.STRICT_STUBS)
  public class OffsetStorageWriterTest {
      private static final String NAMESPACE = "namespace";
      // Connect format - any types should be accepted here
      private static final Map<String, Object> OFFSET_KEY = Collections.singletonMap("key", "key");
      private static final Map<String, Object> OFFSET_VALUE = Collections.singletonMap("key", 12);
  
      // Serialized
      private static final byte[] OFFSET_KEY_SERIALIZED = "key-serialized".getBytes();
      private static final byte[] OFFSET_VALUE_SERIALIZED = "value-serialized".getBytes();
  
      private static final Exception EXCEPTION = new RuntimeException("error");
  
      private final OffsetBackingStore store = mock(OffsetBackingStore.class);
      private final Converter keyConverter = mock(Converter.class);
      private final Converter valueConverter = mock(Converter.class);
      private OffsetStorageWriter writer;
  
      private ExecutorService service;
  
      @BeforeEach
      public void setup() {
          writer = new OffsetStorageWriter(store, NAMESPACE, keyConverter, valueConverter);
          service = Executors.newFixedThreadPool(1);
      }
  
      @AfterEach
      public void teardown() {
          service.shutdownNow();
      }
  
      @Test
      public void testWriteFlush() throws Exception {
          @SuppressWarnings("unchecked")
          Callback<Void> callback = mock(Callback.class);
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, false, null);
  
          writer.offset(OFFSET_KEY, OFFSET_VALUE);
  
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
          verify(callback).onCompletion(isNull(), isNull());
      }
  
      // It should be possible to set offset values to null
      @Test
      public void testWriteNullValueFlush() throws Exception {
          @SuppressWarnings("unchecked")
          Callback<Void> callback = mock(Callback.class);
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, null, null, false, null);
  
          writer.offset(OFFSET_KEY, null);
  
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
          verify(callback).onCompletion(isNull(), isNull());
      }
  
      // It should be possible to use null keys. These aren't actually stored as null since the key is wrapped to include
      // info about the namespace (connector)
      @Test
      public void testWriteNullKeyFlush() throws Exception {
          @SuppressWarnings("unchecked")
          Callback<Void> callback = mock(Callback.class);
          expectStore(null, null, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, false, null);
  
          writer.offset(null, OFFSET_VALUE);
  
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
          verify(callback).onCompletion(isNull(), isNull());
      }
  
      @Test
      public void testNoOffsetsToFlush() throws InterruptedException, TimeoutException {
          assertFalse(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          assertFalse(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
  
          // If no offsets are flushed, we should finish immediately and not have made any calls to the
          // underlying storage layer
          verifyNoInteractions(store);
      }
  
      @Test
      public void testFlushFailureReplacesOffsets() throws Exception {
          // When a flush fails, we shouldn't just lose the offsets. Instead, they should be restored
          // such that a subsequent flush will write them.
  
          @SuppressWarnings("unchecked")
          final Callback<Void> callback = mock(Callback.class);
          // First time the write fails
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, true, null);
          writer.offset(OFFSET_KEY, OFFSET_VALUE);
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
          verify(callback).onCompletion(eq(EXCEPTION), isNull());
  
          // Second time it succeeds
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, false, null);
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback).get(1000, TimeUnit.MILLISECONDS);
          verify(callback).onCompletion(isNull(), isNull());
  
          // Third time it has no data to flush so we won't get past beginFlush()
          assertFalse(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
      }
  
      @Test
      public void testAlreadyFlushing() throws InterruptedException, TimeoutException {
          @SuppressWarnings("unchecked")
          final Callback<Void> callback = mock(Callback.class);
          // Trigger the send, but don't invoke the callback so we'll still be mid-flush
          CountDownLatch allowStoreCompleteCountdown = new CountDownLatch(1);
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, false, allowStoreCompleteCountdown);
  
          writer.offset(OFFSET_KEY, OFFSET_VALUE);
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          assertThrows(TimeoutException.class, () -> writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.doFlush(callback);
          assertThrows(TimeoutException.class, () -> writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          allowStoreCompleteCountdown.countDown();
          assertFalse(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
      }
  
      @Test
      public void testCancelBeforeAwaitFlush() throws InterruptedException, TimeoutException {
          writer.offset(OFFSET_KEY, OFFSET_VALUE);
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          writer.cancelFlush();
      }
  
      @Test
      public void testCancelAfterAwaitFlush() throws Exception {
          @SuppressWarnings("unchecked")
          Callback<Void> callback = mock(Callback.class);
          CountDownLatch allowStoreCompleteCountdown = new CountDownLatch(1);
          // In this test, the write should be cancelled so the callback will not be invoked and is not
          // passed to the expectStore call
          expectStore(OFFSET_KEY, OFFSET_KEY_SERIALIZED, OFFSET_VALUE, OFFSET_VALUE_SERIALIZED, false, allowStoreCompleteCountdown);
  
          writer.offset(OFFSET_KEY, OFFSET_VALUE);
          assertTrue(writer.beginFlush(1000L, TimeUnit.MILLISECONDS));
          // Start the flush, then immediately cancel before allowing the mocked store request to finish
          Future<Void> flushFuture = writer.doFlush(callback);
          writer.cancelFlush();
          allowStoreCompleteCountdown.countDown();
          flushFuture.get(1000, TimeUnit.MILLISECONDS);
      }
  
      /**
       * Expect a request to store data to the underlying OffsetBackingStore.
       *
       * @param key the key for the offset
       * @param keySerialized serialized version of the key
       * @param value the value for the offset
       * @param valueSerialized serialized version of the value
       * @param fail if true, treat
       * @param waitForCompletion if non-null, a CountDownLatch that should be awaited on before
       *                          invoking the callback. A (generous) timeout is still imposed to
       *                          ensure tests complete.
       */
      @SuppressWarnings("unchecked")
      private void expectStore(Map<String, Object> key, byte[] keySerialized,
                               Map<String, Object> value, byte[] valueSerialized,
                               final boolean fail,
                               final CountDownLatch waitForCompletion) {
          List<Object> keyWrapped = Arrays.asList(NAMESPACE, key);
          when(keyConverter.fromConnectData(NAMESPACE, null, keyWrapped)).thenReturn(keySerialized);
          when(valueConverter.fromConnectData(NAMESPACE, null, value)).thenReturn(valueSerialized);
  
          final ArgumentCaptor<Callback<Void>> storeCallback = ArgumentCaptor.forClass(Callback.class);
          final Map<ByteBuffer, ByteBuffer> offsetsSerialized = Collections.singletonMap(
                  keySerialized == null ? null : ByteBuffer.wrap(keySerialized),
                  valueSerialized == null ? null : ByteBuffer.wrap(valueSerialized));
          when(store.set(eq(offsetsSerialized), storeCallback.capture())).thenAnswer(invocation -> {
              final Callback<Void> cb = invocation.getArgument(1);
              return service.submit(() -> {
                  if (waitForCompletion != null)
                      assertTrue(waitForCompletion.await(10000, TimeUnit.MILLISECONDS));
  
                  if (fail) {
                      cb.onCompletion(EXCEPTION, null);
                  } else {
                      cb.onCompletion(null, null);
                  }
                  return null;
              });
          });
      }
  
  }
  