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 com.datastax.oss.driver.internal.core.cql;
  
  import static com.datastax.oss.driver.Assertions.assertThat;
  import static com.datastax.oss.driver.Assertions.assertThatStage;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.anyLong;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.verifyNoMoreInteractions;
  import static org.mockito.Mockito.when;
  
  import com.datastax.oss.driver.api.core.CqlSession;
  import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
  import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
  import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
  import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
  import com.datastax.oss.driver.api.core.cql.ColumnDefinitions;
  import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
  import com.datastax.oss.driver.api.core.cql.QueryTrace;
  import com.datastax.oss.driver.api.core.cql.Row;
  import com.datastax.oss.driver.api.core.cql.SimpleStatement;
  import com.datastax.oss.driver.api.core.cql.TraceEvent;
  import com.datastax.oss.driver.api.core.uuid.Uuids;
  import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
  import com.datastax.oss.driver.internal.core.context.NettyOptions;
  import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
  import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
  import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
  import com.datastax.oss.protocol.internal.util.Bytes;
  import io.netty.util.concurrent.EventExecutor;
  import io.netty.util.concurrent.EventExecutorGroup;
  import java.net.InetAddress;
  import java.nio.ByteBuffer;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.UUID;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  import java.util.concurrent.TimeUnit;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.mockito.ArgumentCaptor;
  import org.mockito.Captor;
  import org.mockito.Mock;
  import org.mockito.junit.MockitoJUnitRunner;
  
  @RunWith(MockitoJUnitRunner.class)
  public class QueryTraceFetcherTest {
  
    private static final UUID TRACING_ID = UUID.randomUUID();
    private static final ByteBuffer PAGING_STATE = Bytes.fromHexString("0xdeadbeef");
    private static final int PORT = 7000;
  
    @Mock private CqlSession session;
    @Mock private InternalDriverContext context;
    @Mock private DriverExecutionProfile config;
    @Mock private DriverExecutionProfile traceConfig;
    @Mock private NettyOptions nettyOptions;
    @Mock private EventExecutorGroup adminEventExecutorGroup;
    @Mock private EventExecutor eventExecutor;
    private InetAddress address = InetAddress.getLoopbackAddress();
  
    @Captor private ArgumentCaptor<SimpleStatement> statementCaptor;
  
    @Before
    public void setup() {
      when(context.getNettyOptions()).thenReturn(nettyOptions);
      when(nettyOptions.adminEventExecutorGroup()).thenReturn(adminEventExecutorGroup);
      when(adminEventExecutorGroup.next()).thenReturn(eventExecutor);
      // Always execute scheduled tasks immediately:
      when(eventExecutor.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)))
          .thenAnswer(
              invocation -> {
                Runnable runnable = invocation.getArgument(0);
                runnable.run();
                // OK because the production code doesn't use the result:
                return null;
              });
  
      when(config.getInt(DefaultDriverOption.REQUEST_TRACE_ATTEMPTS)).thenReturn(3);
      // Doesn't really matter since we mock the scheduler
      when(config.getDuration(DefaultDriverOption.REQUEST_TRACE_INTERVAL)).thenReturn(Duration.ZERO);
      when(config.getString(DefaultDriverOption.REQUEST_CONSISTENCY))
          .thenReturn(DefaultConsistencyLevel.LOCAL_ONE.name());
      when(config.getString(DefaultDriverOption.REQUEST_TRACE_CONSISTENCY))
          .thenReturn(DefaultConsistencyLevel.ONE.name());
  
      when(config.withString(
              DefaultDriverOption.REQUEST_CONSISTENCY, DefaultConsistencyLevel.ONE.name()))
          .thenReturn(traceConfig);
    }
  
    @Test
    public void should_succeed_when_both_queries_succeed_immediately() {
      // Given
      CompletionStage<AsyncResultSet> sessionRow = completeSessionRow();
      CompletionStage<AsyncResultSet> eventRows = singlePageEventRows();
      when(session.executeAsync(any(SimpleStatement.class)))
          .thenAnswer(invocation -> sessionRow)
          .thenAnswer(invocation -> eventRows);
  
      // When
      QueryTraceFetcher fetcher = new QueryTraceFetcher(TRACING_ID, session, context, config);
      CompletionStage<QueryTrace> traceFuture = fetcher.fetch();
  
      // Then
      verify(session, times(2)).executeAsync(statementCaptor.capture());
      List<SimpleStatement> statements = statementCaptor.getAllValues();
      assertSessionQuery(statements.get(0));
      SimpleStatement statement = statements.get(1);
      assertEventsQuery(statement);
      verifyNoMoreInteractions(session);
  
      assertThatStage(traceFuture)
          .isSuccess(
              trace -> {
                assertThat(trace.getTracingId()).isEqualTo(TRACING_ID);
                assertThat(trace.getRequestType()).isEqualTo("mock request");
                assertThat(trace.getDurationMicros()).isEqualTo(42);
                assertThat(trace.getCoordinatorAddress().getAddress()).isEqualTo(address);
                assertThat(trace.getCoordinatorAddress().getPort()).isEqualTo(PORT);
                assertThat(trace.getParameters())
                    .hasSize(2)
                    .containsEntry("key1", "value1")
                    .containsEntry("key2", "value2");
                assertThat(trace.getStartedAt()).isEqualTo(0);
  
                List<TraceEvent> events = trace.getEvents();
                assertThat(events).hasSize(3);
                for (int i = 0; i < events.size(); i++) {
                  TraceEvent event = events.get(i);
                  assertThat(event.getActivity()).isEqualTo("mock activity " + i);
                  assertThat(event.getTimestamp()).isEqualTo(i);
                  assertThat(event.getSourceAddress()).isNotNull();
                  assertThat(event.getSourceAddress().getAddress()).isEqualTo(address);
                  assertThat(event.getSourceAddress().getPort()).isEqualTo(PORT);
                  assertThat(event.getSourceElapsedMicros()).isEqualTo(i);
                  assertThat(event.getThreadName()).isEqualTo("mock thread " + i);
                }
              });
    }
  
    /**
     * This should not happen with a sane configuration, but we need to handle it in case {@link
     * DefaultDriverOption#REQUEST_PAGE_SIZE} is set ridiculously low.
     */
    @Test
    public void should_succeed_when_events_query_is_paged() {
      // Given
      CompletionStage<AsyncResultSet> sessionRow = completeSessionRow();
      CompletionStage<AsyncResultSet> eventRows1 = multiPageEventRows1();
      CompletionStage<AsyncResultSet> eventRows2 = multiPageEventRows2();
      when(session.executeAsync(any(SimpleStatement.class)))
          .thenAnswer(invocation -> sessionRow)
          .thenAnswer(invocation -> eventRows1)
          .thenAnswer(invocation -> eventRows2);
  
      // When
      QueryTraceFetcher fetcher = new QueryTraceFetcher(TRACING_ID, session, context, config);
      CompletionStage<QueryTrace> traceFuture = fetcher.fetch();
  
      // Then
      verify(session, times(3)).executeAsync(statementCaptor.capture());
      List<SimpleStatement> statements = statementCaptor.getAllValues();
      assertSessionQuery(statements.get(0));
      assertEventsQuery(statements.get(1));
      assertEventsQuery(statements.get(2));
      assertThat(statements.get(2).getPagingState()).isEqualTo(PAGING_STATE);
      verifyNoMoreInteractions(session);
  
      assertThatStage(traceFuture).isSuccess(trace -> assertThat(trace.getEvents()).hasSize(2));
    }
  
    @Test
    public void should_retry_when_session_row_is_incomplete() {
      // Given
      CompletionStage<AsyncResultSet> sessionRow1 = incompleteSessionRow();
      CompletionStage<AsyncResultSet> sessionRow2 = completeSessionRow();
      CompletionStage<AsyncResultSet> eventRows = singlePageEventRows();
      when(session.executeAsync(any(SimpleStatement.class)))
          .thenAnswer(invocation -> sessionRow1)
          .thenAnswer(invocation -> sessionRow2)
          .thenAnswer(invocation -> eventRows);
  
      // When
      QueryTraceFetcher fetcher = new QueryTraceFetcher(TRACING_ID, session, context, config);
      CompletionStage<QueryTrace> traceFuture = fetcher.fetch();
  
      // Then
      verify(session, times(3)).executeAsync(statementCaptor.capture());
      List<SimpleStatement> statements = statementCaptor.getAllValues();
      assertSessionQuery(statements.get(0));
      assertSessionQuery(statements.get(1));
      assertEventsQuery(statements.get(2));
      verifyNoMoreInteractions(session);
  
      assertThatStage(traceFuture)
          .isSuccess(
              trace -> {
                assertThat(trace.getTracingId()).isEqualTo(TRACING_ID);
                assertThat(trace.getRequestType()).isEqualTo("mock request");
                assertThat(trace.getDurationMicros()).isEqualTo(42);
                assertThat(trace.getCoordinatorAddress().getAddress()).isEqualTo(address);
                assertThat(trace.getCoordinatorAddress().getPort()).isEqualTo(PORT);
                assertThat(trace.getParameters())
                    .hasSize(2)
                    .containsEntry("key1", "value1")
                    .containsEntry("key2", "value2");
                assertThat(trace.getStartedAt()).isEqualTo(0);
  
                List<TraceEvent> events = trace.getEvents();
                assertThat(events).hasSize(3);
                for (int i = 0; i < events.size(); i++) {
                  TraceEvent event = events.get(i);
                  assertThat(event.getActivity()).isEqualTo("mock activity " + i);
                  assertThat(event.getTimestamp()).isEqualTo(i);
                  assertThat(event.getSourceAddress()).isNotNull();
                  assertThat(event.getSourceAddress().getAddress()).isEqualTo(address);
                  assertThat(event.getSourceAddress().getPort()).isEqualTo(PORT);
                  assertThat(event.getSourceElapsedMicros()).isEqualTo(i);
                  assertThat(event.getThreadName()).isEqualTo("mock thread " + i);
                }
              });
    }
  
    @Test
    public void should_fail_when_session_query_fails() {
      // Given
      RuntimeException mockError = new RuntimeException("mock error");
      when(session.executeAsync(any(SimpleStatement.class)))
          .thenReturn(CompletableFutures.failedFuture(mockError));
  
      // When
      QueryTraceFetcher fetcher = new QueryTraceFetcher(TRACING_ID, session, context, config);
      CompletionStage<QueryTrace> traceFuture = fetcher.fetch();
  
      // Then
      verify(session).executeAsync(statementCaptor.capture());
      SimpleStatement statement = statementCaptor.getValue();
      assertSessionQuery(statement);
      verifyNoMoreInteractions(session);
  
      assertThatStage(traceFuture).isFailed(error -> assertThat(error).isSameAs(mockError));
    }
  
    @Test
    public void should_fail_when_session_query_still_incomplete_after_max_tries() {
      // Given
      CompletionStage<AsyncResultSet> sessionRow1 = incompleteSessionRow();
      CompletionStage<AsyncResultSet> sessionRow2 = incompleteSessionRow();
      CompletionStage<AsyncResultSet> sessionRow3 = incompleteSessionRow();
      when(session.executeAsync(any(SimpleStatement.class)))
          .thenAnswer(invocation -> sessionRow1)
          .thenAnswer(invocation -> sessionRow2)
          .thenAnswer(invocation -> sessionRow3);
  
      // When
      QueryTraceFetcher fetcher = new QueryTraceFetcher(TRACING_ID, session, context, config);
      CompletionStage<QueryTrace> traceFuture = fetcher.fetch();
  
      // Then
      verify(session, times(3)).executeAsync(statementCaptor.capture());
      List<SimpleStatement> statements = statementCaptor.getAllValues();
      for (int i = 0; i < 3; i++) {
        assertSessionQuery(statements.get(i));
      }
  
      assertThatStage(traceFuture)
          .isFailed(
              error ->
                  assertThat(error.getMessage())
                      .isEqualTo(
                          String.format("Trace %s still not complete after 3 attempts", TRACING_ID)));
    }
  
    private CompletionStage<AsyncResultSet> completeSessionRow() {
      return sessionRow(42);
    }
  
    private CompletionStage<AsyncResultSet> incompleteSessionRow() {
      return sessionRow(null);
    }
  
    private CompletionStage<AsyncResultSet> sessionRow(Integer duration) {
      Row row = mock(Row.class);
      ColumnDefinitions definitions = mock(ColumnDefinitions.class);
      when(row.getColumnDefinitions()).thenReturn(definitions);
      when(row.getString("request")).thenReturn("mock request");
      if (duration == null) {
        when(row.isNull("duration")).thenReturn(true);
      } else {
        when(row.getInt("duration")).thenReturn(duration);
      }
      when(row.getInetAddress("coordinator")).thenReturn(address);
      when(definitions.contains("coordinator_port")).thenReturn(true);
      when(row.getInt("coordinator_port")).thenReturn(PORT);
      when(row.getMap("parameters", String.class, String.class))
          .thenReturn(ImmutableMap.of("key1", "value1", "key2", "value2"));
      when(row.isNull("started_at")).thenReturn(false);
      when(row.getInstant("started_at")).thenReturn(Instant.EPOCH);
  
      AsyncResultSet rs = mock(AsyncResultSet.class);
      when(rs.one()).thenReturn(row);
      return CompletableFuture.completedFuture(rs);
    }
  
    private CompletionStage<AsyncResultSet> singlePageEventRows() {
      List<Row> rows = new ArrayList<>();
      for (int i = 0; i < 3; i++) {
        rows.add(eventRow(i));
      }
  
      AsyncResultSet rs = mock(AsyncResultSet.class);
      when(rs.currentPage()).thenReturn(rows);
  
      ExecutionInfo executionInfo = mock(ExecutionInfo.class);
      when(executionInfo.getPagingState()).thenReturn(null);
      when(rs.getExecutionInfo()).thenReturn(executionInfo);
  
      return CompletableFuture.completedFuture(rs);
    }
  
    private CompletionStage<AsyncResultSet> multiPageEventRows1() {
      AsyncResultSet rs = mock(AsyncResultSet.class);
  
      ImmutableList<Row> rows = ImmutableList.of(eventRow(0));
      when(rs.currentPage()).thenReturn(rows);
  
      ExecutionInfo executionInfo = mock(ExecutionInfo.class);
      when(executionInfo.getPagingState()).thenReturn(PAGING_STATE);
      when(rs.getExecutionInfo()).thenReturn(executionInfo);
  
      return CompletableFuture.completedFuture(rs);
    }
  
    private CompletionStage<AsyncResultSet> multiPageEventRows2() {
      AsyncResultSet rs = mock(AsyncResultSet.class);
  
      ImmutableList<Row> rows = ImmutableList.of(eventRow(1));
      when(rs.currentPage()).thenReturn(rows);
  
      ExecutionInfo executionInfo = mock(ExecutionInfo.class);
      when(executionInfo.getPagingState()).thenReturn(null);
      when(rs.getExecutionInfo()).thenReturn(executionInfo);
  
      return CompletableFuture.completedFuture(rs);
    }
  
    private Row eventRow(int i) {
      Row row = mock(Row.class);
      ColumnDefinitions definitions = mock(ColumnDefinitions.class);
      when(row.getColumnDefinitions()).thenReturn(definitions);
      when(row.getString("activity")).thenReturn("mock activity " + i);
      when(row.getUuid("event_id")).thenReturn(Uuids.startOf(i));
      when(row.getInetAddress("source")).thenReturn(address);
      when(definitions.contains("source_port")).thenReturn(true);
      when(row.getInt("source_port")).thenReturn(PORT);
      when(row.getInt("source_elapsed")).thenReturn(i);
      when(row.getString("thread")).thenReturn("mock thread " + i);
      return row;
    }
  
    private void assertSessionQuery(SimpleStatement statement) {
      assertThat(statement.getQuery())
          .isEqualTo("SELECT * FROM system_traces.sessions WHERE session_id = ?");
      assertThat(statement.getPositionalValues()).containsOnly(TRACING_ID);
      assertThat(statement.getExecutionProfile()).isEqualTo(traceConfig);
    }
  
    private void assertEventsQuery(SimpleStatement statement) {
      assertThat(statement.getQuery())
          .isEqualTo("SELECT * FROM system_traces.events WHERE session_id = ?");
      assertThat(statement.getPositionalValues()).containsOnly(TRACING_ID);
      assertThat(statement.getExecutionProfile()).isEqualTo(traceConfig);
    }
  }
  