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.activemq.artemis.cli.commands.messages.perf;
  
  import javax.jms.BytesMessage;
  import javax.jms.CompletionListener;
  import javax.jms.JMSException;
  import javax.jms.Message;
  
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.function.BooleanSupplier;
  
  import io.netty.util.concurrent.OrderedEventExecutor;
  import org.HdrHistogram.SingleWriterRecorder;
  import org.apache.activemq.artemis.cli.commands.messages.perf.AsyncJms2ProducerFacade.SendAttemptResult;
  
  public abstract class SkeletalProducerLoadGenerator implements CompletionListener, ProducerLoadGenerator {
  
     protected final AsyncJms2ProducerFacade producer;
     private final OrderedEventExecutor executor;
     protected final BooleanSupplier keepOnSending;
     protected final MicrosTimeProvider timeProvider;
     private final String group;
     private final byte[] messageContent;
     private BytesMessage messageToSend;
     protected boolean closed;
     protected volatile boolean stopLoad;
     private final SingleWriterRecorder waitLatencies;
     private final SingleWriterRecorder sendCompletedLatencies;
     private final AtomicLong unprocessedCompletions;
     private final AtomicBoolean scheduledProcessingCompletions;
     private volatile Exception fatalException;
     private boolean stopHandlingCompletions;
  
     public SkeletalProducerLoadGenerator(final AsyncJms2ProducerFacade producer,
                                          final OrderedEventExecutor executor,
                                          final MicrosTimeProvider timeProvider,
                                          final BooleanSupplier keepOnSending,
                                          final String group,
                                          final byte[] msgContent,
                                          final SingleWriterRecorder sendCompletedLatencies,
                                          final SingleWriterRecorder waitLatencies) {
        this.sendCompletedLatencies = sendCompletedLatencies;
        this.waitLatencies = waitLatencies;
        this.producer = producer;
        this.executor = executor;
        this.timeProvider = timeProvider;
        this.keepOnSending = keepOnSending;
        this.group = group;
        this.messageContent = msgContent;
        this.messageToSend = null;
        this.closed = false;
        this.stopLoad = false;
        this.unprocessedCompletions = new AtomicLong();
        this.scheduledProcessingCompletions = new AtomicBoolean();
        this.fatalException = null;
        this.stopHandlingCompletions = false;
     }
  
     @Override
     public Exception getFatalException() {
        return fatalException;
     }
  
     @Override
     public SingleWriterRecorder getSendCompletedLatencies() {
        return sendCompletedLatencies;
     }
  
     @Override
     public SingleWriterRecorder getWaitLatencies() {
        return waitLatencies;
     }
  
     @Override
     public AsyncJms2ProducerFacade getProducer() {
        return producer;
     }
  
     @Override
     public boolean isCompleted() {
        if (stopLoad && fatalException != null) {
           return true;
        }
        return stopLoad && producer.getMessageCompleted() == producer.getMessageSent();
     }
  
     @Override
     public OrderedEventExecutor getExecutor() {
        return executor;
     }
  
     protected final void asyncContinue() {
        asyncContinue(0);
     }
  
     protected final void asyncContinue(final long usDelay) {
        if (usDelay == 0) {
           executor.execute(this);
        } else {
           executor.schedule(this, usDelay, TimeUnit.MICROSECONDS);
        }
     }
  
     protected final boolean trySend(final long sendTime) {
        return trySend(sendTime, sendTime);
     }
  
     protected final boolean trySend(final long expectedSendTime, final long sendTime) {
        assert executor.inEventLoop();
        assert !closed;
        try {
           if (messageToSend == null) {
              messageToSend = producer.createBytesMessage();
              messageToSend.writeBytes(this.messageContent);
           }
           messageToSend.setLongProperty("time", sendTime);
           if (group != null) {
              messageToSend.setStringProperty("JMSXGroupID", group);
           }
           final SendAttemptResult result = producer.trySend(messageToSend, this, this);
           if (result != SendAttemptResult.NotAvailable) {
              messageToSend = null;
              if (result == SendAttemptResult.Success) {
                 if (waitLatencies != null) {
                    waitLatencies.recordValue(sendTime - expectedSendTime);
                 }
              }
           }
           return result == SendAttemptResult.Success;
        } catch (final JMSException e) {
           onSendErrored(e);
           return false;
        }
     }
  
     @Override
     public void onCompletion(final Message message) {
        asyncOnSendCompleted(message, null);
     }
  
     @Override
     public void onException(final Message message, final Exception exception) {
        asyncOnSendCompleted(message, exception);
     }
  
     private void asyncOnSendCompleted(final Message message, Exception completionError) {
        if (stopHandlingCompletions) {
           return;
        }
        if (completionError == null) {
           try {
              recordSendCompletionLatency(message);
              unprocessedCompletions.incrementAndGet();
              scheduleProcessingCompletions();
           } catch (final JMSException jmsException) {
              completionError = jmsException;
           }
        }
        if (completionError != null) {
           stopHandlingCompletions = true;
           final Exception fatal = completionError;
           executor.execute(() -> onSendErrored(fatal));
        }
     }
  
     private void onSendErrored(final Exception fatal) {
        assert executor.inEventLoop();
        if (fatalException != null) {
           return;
        }
        producer.onSendErrored();
        fatalException = fatal;
        stopLoad = true;
        closed = true;
     }
  
     private void scheduleProcessingCompletions() {
        if (unprocessedCompletions.get() > 0 && scheduledProcessingCompletions.compareAndSet(false, true)) {
           executor.execute(this::processCompletions);
        }
     }
  
     private void processCompletions() {
        assert executor.inEventLoop();
        assert scheduledProcessingCompletions.get();
        if (fatalException != null) {
           return;
        }
        final long completions = unprocessedCompletions.getAndSet(0);
        for (long i = 0; i < completions; i++) {
           final JMSException completionException = producer.onSendCompleted();
           if (completionException != null) {
              fatalException = completionException;
              return;
           }
        }
        scheduledProcessingCompletions.set(false);
        scheduleProcessingCompletions();
     }
  
     private void recordSendCompletionLatency(final Message message) throws JMSException {
        final long time = message.getLongProperty("time");
        final long elapsedMicros = timeProvider.now() - time;
        sendCompletedLatencies.recordValue(elapsedMicros);
     }
  
     @Override
     public Future<?> asyncClose(final Runnable onClosed) {
        return executor.submit(() -> onClose(onClosed));
     }
  
     private void onClose(final Runnable onClosed) {
        assert executor.inEventLoop();
        if (closed) {
           onClosed.run();
           return;
        }
        closed = true;
        // no need for this anymore
        messageToSend = null;
        producer.requestClose(onClosed);
     }
  }
  