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.Destination;
  import javax.jms.JMSException;
  import javax.jms.Message;
  import javax.jms.MessageProducer;
  import javax.jms.Session;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.atomic.AtomicLongFieldUpdater;
  
  import static java.util.Objects.requireNonNull;
  
  public final class AsyncJms2ProducerFacade {
  
     private final long id;
     protected final Session session;
     private final MessageProducer producer;
  
     /*
      * maxPending limits the number of in-flight sent messages
      * in a way that if the limit is reached and a single completion arrive,
      * a subsequent send attempt will succeed.
      */
     private long pending;
     private final long maxPending;
  
     /*
      * Tracking sent messages in transaction requires using 2 separate counters
      * ie pendingMsgInTransaction, completedMsgInTransaction
      * because, using just one won't allow tracking completions of previously sent messages in order to commit
      * the transaction while there are no more in-flight ones.
      */
     private final long transactionCapacity;
     private long pendingMsgInTransaction;
     private long completedMsgInTransaction;
  
     private final List<Runnable> availableObservers;
     private final List<Runnable> closedObservers;
  
     private static final AtomicLongFieldUpdater<AsyncJms2ProducerFacade> MESSAGE_SENT_UPDATER = AtomicLongFieldUpdater.newUpdater(AsyncJms2ProducerFacade.class, "messageSent");
     private static final AtomicLongFieldUpdater<AsyncJms2ProducerFacade> MESSAGE_COMPLETED_UPDATER = AtomicLongFieldUpdater.newUpdater(AsyncJms2ProducerFacade.class, "messageCompleted");
     private static final AtomicLongFieldUpdater<AsyncJms2ProducerFacade> NOT_AVAILABLE_UPDATER = AtomicLongFieldUpdater.newUpdater(AsyncJms2ProducerFacade.class, "notAvailable");
  
     private volatile long messageSent;
     private volatile long messageCompleted;
     private volatile long notAvailable;
  
     private boolean closing;
     private boolean closed;
     private final Destination destination;
  
     public AsyncJms2ProducerFacade(final long id,
                                    final Session session,
                                    final MessageProducer producer,
                                    final Destination destination,
                                    final long maxPending,
                                    final long transactionCapacity) {
        this.id = id;
        this.session = requireNonNull(session);
        this.producer = requireNonNull(producer);
        this.destination = destination;
        this.pending = 0;
        this.maxPending = transactionCapacity > 0 && maxPending > 0 ? Math.max(maxPending, transactionCapacity) : maxPending;
        this.availableObservers = new ArrayList<>(1);
        this.closedObservers = new ArrayList<>(1);
        this.messageSent = 0;
        this.messageCompleted = 0;
        this.notAvailable = 0;
        try {
           if (transactionCapacity < 0) {
              throw new IllegalStateException("transactionCapacity must be >= 0");
           }
           if (transactionCapacity > 0) {
              if (!session.getTransacted()) {
                 throw new IllegalStateException("session must be transacted with transactionCapacity != 0");
              }
           } else {
              if (session.getTransacted()) {
                 throw new IllegalStateException("session cannot be transacted with transactionCapacity = 0");
              }
           }
        } catch (final JMSException ex) {
           throw new IllegalStateException(ex);
        }
        this.transactionCapacity = transactionCapacity;
        this.pendingMsgInTransaction = 0;
        this.completedMsgInTransaction = 0;
        this.closing = false;
        this.closed = false;
     }
  
     public long getId() {
        return id;
     }
  
     public Destination getDestination() {
        return destination;
     }
  
     BytesMessage createBytesMessage() throws JMSException {
        return session.createBytesMessage();
     }
  
     private void addedPendingSend() {
        if (transactionCapacity > 0 && pendingMsgInTransaction == transactionCapacity) {
           throw new IllegalStateException("reached max in-flight transacted sent messages");
        }
        if (maxPending > 0 && pending == maxPending) {
           throw new IllegalStateException("reached max in-flight sent messages");
        }
        pending++;
        pendingMsgInTransaction++;
     }
  
     /**
      * If {@code true}, a subsequent {@link #trySend} would return {@link SendAttemptResult#Success}. Otherwise, a
      * subsequent {@link #trySend} would return {@link SendAttemptResult#NotAvailable}.
      */
     private boolean isAvailable() {
        if (maxPending > 0 && pending == maxPending) {
           return false;
        }
        return transactionCapacity == 0 || pendingMsgInTransaction != transactionCapacity;
     }
  
     public enum SendAttemptResult {
        Closing, Closed, NotAvailable, Success
     }
  
     public SendAttemptResult trySend(final Message message,
                                      final CompletionListener completionListener,
                                      final Runnable availableObserver) throws JMSException {
        if (closing) {
           return SendAttemptResult.Closing;
        }
        if (closed) {
           return SendAttemptResult.Closed;
        }
        if (!isAvailable()) {
           availableObservers.add(availableObserver);
           orderedIncrementNotAvailable();
           return SendAttemptResult.NotAvailable;
        }
        producer.send(message, completionListener);
        orderedIncrementSent();
        addedPendingSend();
        return SendAttemptResult.Success;
     }
  
     public void onSendErrored() {
        if (closed) {
           return;
        }
        availableObservers.clear();
        closedObservers.forEach(Runnable::run);
        closedObservers.clear();
        closed = true;
     }
  
     public JMSException onSendCompleted() {
        if (closed) {
           return null;
        }
        JMSException completionError = null;
        orderedIncrementCompleted();
        if (transactionCapacity > 0 && completedMsgInTransaction == transactionCapacity) {
           throw new IllegalStateException("cannot complete more send");
        }
        if (pending == 0) {
           throw new IllegalStateException("cannot complete more send");
        }
        pending--;
        completedMsgInTransaction++;
        if (transactionCapacity > 0) {
           if (completedMsgInTransaction == transactionCapacity || (closing && pending == 0)) {
              completedMsgInTransaction = 0;
              pendingMsgInTransaction = 0;
              try {
                 session.commit();
              } catch (final JMSException fatal) {
                 completionError = fatal;
                 closing = true;
              }
              if (closing) {
                 closing = false;
                 closed = true;
                 closedObservers.forEach(Runnable::run);
                 closedObservers.clear();
              } else if (isAvailable()) {
                 availableObservers.forEach(Runnable::run);
                 availableObservers.clear();
              }
           }
        } else {
           if (closing && pending == 0) {
              closing = false;
              closed = true;
              closedObservers.forEach(Runnable::run);
              closedObservers.clear();
           } else if (isAvailable()) {
              availableObservers.forEach(Runnable::run);
              availableObservers.clear();
           }
        }
        return completionError;
     }
  
     public long getMessageSent() {
        return messageSent;
     }
  
     private void orderedIncrementSent() {
        MESSAGE_SENT_UPDATER.lazySet(this, messageSent + 1);
     }
  
     public long getMessageCompleted() {
        return messageCompleted;
     }
  
     private void orderedIncrementCompleted() {
        MESSAGE_COMPLETED_UPDATER.lazySet(this, messageCompleted + 1);
     }
  
     public long getNotAvailable() {
        return notAvailable;
     }
  
     private void orderedIncrementNotAvailable() {
        NOT_AVAILABLE_UPDATER.lazySet(this, notAvailable + 1);
     }
  
     public void requestClose() {
        requestClose(() -> {
        });
     }
  
     public void requestClose(final Runnable onClosed) {
        if (closed) {
           onClosed.run();
           return;
        }
        if (closing) {
           closedObservers.add(onClosed);
           return;
        }
        availableObservers.clear();
        if (pending > 0) {
           closing = true;
           closedObservers.add(onClosed);
        } else {
           closed = true;
           onClosed.run();
        }
     }
  }
  