code: /*
   * Copyright (c) 2019 the Eclipse Milo Authors
   *
   * This program and the accompanying materials are made
   * available under the terms of the Eclipse Public License 2.0
   * which is available at https://www.eclipse.org/legal/epl-2.0/
   *
   * SPDX-License-Identifier: EPL-2.0
   */
  
  package org.eclipse.milo.opcua.sdk.client.subscriptions;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.stream.Collectors;
  
  import com.google.common.collect.ArrayListMultimap;
  import com.google.common.collect.ListMultimap;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Multimaps;
  import org.eclipse.milo.opcua.sdk.client.OpcUaClient;
  import org.eclipse.milo.opcua.sdk.client.api.subscriptions.UaMonitoredItem;
  import org.eclipse.milo.opcua.sdk.client.nodes.UaVariableNode;
  import org.eclipse.milo.opcua.sdk.core.util.GroupMapCollate;
  import org.eclipse.milo.opcua.stack.core.AttributeId;
  import org.eclipse.milo.opcua.stack.core.Identifiers;
  import org.eclipse.milo.opcua.stack.core.StatusCodes;
  import org.eclipse.milo.opcua.stack.core.UaException;
  import org.eclipse.milo.opcua.stack.core.types.builtin.StatusCode;
  import org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.UInteger;
  import org.eclipse.milo.opcua.stack.core.types.enumerated.MonitoringMode;
  import org.eclipse.milo.opcua.stack.core.util.FutureUtils;
  import org.jetbrains.annotations.Nullable;
  
  import static org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.Unsigned.uint;
  
  
  /**
   * Manages the batch execution of one or more {@link MonitoringMode} modifications to one or more MonitoredItems.
   * <p>
   * Handles splitting the execution into multiple service calls as needed based on the requested modifications and the
   * operation limits reported by the server.
   */
  public class BatchSetMonitoringMode {
  
      private final Map<OpcUaMonitoredItem, MonitoringMode> monitoringModesByItem =
          Collections.synchronizedMap(new LinkedHashMap<>());
  
      private final ListMultimap<OpcUaMonitoredItem, CompletableFuture<SetMonitoringModeResult>> futuresByItem =
          Multimaps.synchronizedListMultimap(ArrayListMultimap.create());
  
      private final List<CompletableFuture<SetMonitoringModeResult>> resultFutures =
          Collections.synchronizedList(new ArrayList<>());
  
      private final AtomicInteger serviceInvocationCount = new AtomicInteger(0);
  
      private final OpcUaClient client;
      private final OpcUaSubscription subscription;
  
      public BatchSetMonitoringMode(ManagedSubscription subscription) {
          this(subscription.getClient(), subscription.getSubscription());
      }
  
      public BatchSetMonitoringMode(OpcUaClient client, OpcUaSubscription subscription) {
          this.client = client;
          this.subscription = subscription;
      }
  
      /**
       * Get the number of service invocations that were needed to execute this batch.
       *
       * @return the number of service invocations that were needed to execute this batch.
       */
      public int getServiceInvocationCount() {
          return serviceInvocationCount.get();
      }
  
      /**
       * Add an {@link OpcUaMonitoredItem} to set a new {@link MonitoringMode} on as part of a batch.
       *
       * @param monitoredItem  the {@link OpcUaMonitoredItem} to modify.
       * @param monitoringMode the new {@link MonitoringMode}.
       * @return a {@link CompletableFuture} that always completes successfully with a {@link SetMonitoringModeResult}.
       */
      public CompletableFuture<SetMonitoringModeResult> add(
          OpcUaMonitoredItem monitoredItem,
          MonitoringMode monitoringMode
      ) {
  
          monitoringModesByItem.put(monitoredItem, monitoringMode);
  
          CompletableFuture<SetMonitoringModeResult> future = new CompletableFuture<>();
          futuresByItem.put(monitoredItem, future);
          resultFutures.add(future);
          return future;
      }
  
      /**
       * Execute this batch operation and return a List of {@link SetMonitoringModeResult}s the same size and order as
       * calls to {@link #add(OpcUaMonitoredItem, MonitoringMode)} were made.
       *
       * @return a List of {@link SetMonitoringModeResult}s the same size and order as calls to
       * {@link #add(OpcUaMonitoredItem, MonitoringMode)} were made.
       */
      public List<SetMonitoringModeResult> execute() throws InterruptedException {
          try {
              return executeAsync().get();
          } catch (ExecutionException e) {
              throw new IllegalStateException(e);
          }
      }
  
      /**
       * Execute this batch operation and return a {@link CompletableFuture} with a List of
       * {@link SetMonitoringModeResult}s the same size and order as calls to
       * {@link #add(OpcUaMonitoredItem, MonitoringMode)} were made.
       * <p>
       * The returned CompletableFuture always completes successfully.
       *
       * @return a {@link CompletableFuture} with a List of {@link SetMonitoringModeResult}s the same size and order as
       * calls to {@link #add(OpcUaMonitoredItem, MonitoringMode)} were made.
       */
      public CompletableFuture<List<SetMonitoringModeResult>> executeAsync() {
          return readOperationLimit(client).thenCompose(this::executeAsync);
      }
  
      private CompletableFuture<List<SetMonitoringModeResult>> executeAsync(UInteger operationLimit) {
          List<Map.Entry<OpcUaMonitoredItem, MonitoringMode>> entries =
              new ArrayList<>(monitoringModesByItem.entrySet());
  
          CompletableFuture<List<SetMonitoringModeResult>> resultsFuture = GroupMapCollate.groupMapCollate(
              entries,
              Map.Entry::getValue,
              monitoringMode -> itemsAndModes -> {
                  List<CompletableFuture<List<SetMonitoringModeResult>>> partitionFutures =
                      Lists.partition(itemsAndModes, operationLimit.intValue())
                          .stream()
                          .map(this::setMonitoringModeAsync)
                          .collect(Collectors.toList());
  
                  return FutureUtils.flatSequence(partitionFutures);
              }
          );
  
          return resultsFuture.thenCompose(results -> {
              List<OpcUaMonitoredItem> items =
                  new ArrayList<>(monitoringModesByItem.keySet());
  
              assert items.size() == results.size();
  
              for (int i = 0; i < items.size(); i++) {
                  OpcUaMonitoredItem item = items.get(i);
                  SetMonitoringModeResult result = results.get(i);
  
                  List<CompletableFuture<SetMonitoringModeResult>> futures;
                  synchronized (futuresByItem) {
                      futures = new ArrayList<>(futuresByItem.get(item));
                  }
                  futures.forEach(f -> f.complete(result));
              }
  
              return FutureUtils.sequence(resultFutures);
          });
      }
  
      private CompletableFuture<List<SetMonitoringModeResult>> setMonitoringModeAsync(
          List<Map.Entry<OpcUaMonitoredItem, MonitoringMode>> itemsAndModes
      ) {
  
          serviceInvocationCount.incrementAndGet();
  
          MonitoringMode monitoringMode = itemsAndModes.get(0).getValue();
  
          List<UaMonitoredItem> items = itemsAndModes.stream()
              .map(Map.Entry::getKey)
              .collect(Collectors.toList());
  
          CompletableFuture<List<SetMonitoringModeResult>> resultsFuture =
              subscription.setMonitoringMode(monitoringMode, items).thenApply(
                  statusCodes ->
                      statusCodes.stream()
                          .map(statusCode -> new SetMonitoringModeResult(StatusCode.GOOD, statusCode))
                          .collect(Collectors.toList())
              );
  
          return resultsFuture.exceptionally(ex -> {
              StatusCode serviceResult = UaException
                  .extractStatusCode(ex)
                  .orElse(new StatusCode(StatusCodes.Bad_UnexpectedError));
  
              SetMonitoringModeResult result = new SetMonitoringModeResult(serviceResult);
  
              return Collections.nCopies(items.size(), result);
          });
      }
  
      private static CompletableFuture<UInteger> readOperationLimit(OpcUaClient client) {
          CompletableFuture<UaVariableNode> nodeFuture = client.getAddressSpace().getVariableNodeAsync(
              Identifiers.Server_ServerCapabilities_OperationLimits_MaxMonitoredItemsPerCall
          );
  
          return nodeFuture.thenCompose(
              variableNode ->
                  variableNode.readAttributeAsync(AttributeId.Value)
                      .thenApply(v -> (UInteger) v.getValue().getValue())
                      .exceptionally(ex -> uint(1000))
          );
      }
  
      /**
       * Holds the operation- and service-level results of an operation that is part of a
       * {@link BatchSetMonitoringMode} execution.
       */
      public static class SetMonitoringModeResult implements BatchItemResult<StatusCode> {
  
          private final StatusCode serviceResult;
          private final StatusCode operationResult;
  
          SetMonitoringModeResult(StatusCode serviceResult) {
              this(serviceResult, null);
          }
  
          SetMonitoringModeResult(StatusCode serviceResult, @Nullable StatusCode operationResult) {
              this.serviceResult = serviceResult;
              this.operationResult = operationResult;
          }
  
          @Override
          public StatusCode serviceResult() {
              return serviceResult;
          }
  
          @Override
          public Optional<StatusCode> operationResult() {
              return Optional.ofNullable(operationResult);
          }
  
          public boolean isServiceResultGood() {
              return serviceResult.isGood();
          }
  
          public boolean isOperationResultGood() {
              return operationResult().map(StatusCode::isGood).orElse(false);
          }
  
      }
  
  }
  