[Pipe] Batch OPC UA sink writes

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/opcua/client/IoTDBOpcUaClient.java

@@ -61,8 +61,10 @@
 
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
 import static org.apache.iotdb.db.pipe.sink.protocol.opcua.server.OpcUaNameSpace.convertToOpcDataType;
@@ -132,6 +134,7 @@ private void transferTabletRowForClientServerModel(
     Object value = null;
     long timestamp = 0;
     NodeId opcDataType = null;
+    final List<OpcUaWriteRequest> writeRequests = new ArrayList<>();
 
     for (int i = 0; i < measurementSchemas.size(); ++i) {
       if (Objects.isNull(values.get(i))) {
@@ -156,67 +159,149 @@ private void transferTabletRowForClientServerModel(
 
       final long utcTimestamp = timestampToUtc(timestamps.get(timestamps.size() > 1 ? i : 0));
       if (Objects.isNull(sink.getValueName())) {
-        writeValue(
-            values.get(i),
-            utcTimestamp,
-            convertToOpcDataType(type),
-            currentQuality,
-            segments,
-            name);
+        writeRequests.add(
+            new OpcUaWriteRequest(
+                values.get(i),
+                utcTimestamp,
+                convertToOpcDataType(type),
+                currentQuality,
+                segments,
+                name));
       } else {
         value = values.get(i);
         timestamp = utcTimestamp;
         opcDataType = convertToOpcDataType(type);
       }
     }
-    if (Objects.isNull(value)) {
+    if (Objects.nonNull(value)) {
+      writeRequests.add(
+          new OpcUaWriteRequest(value, timestamp, opcDataType, currentQuality, segments, null));
+    }
+
+    writeValues(writeRequests);
+  }
+
+  private void writeValues(final List<OpcUaWriteRequest> writeRequests) throws Exception {
+    if (writeRequests.isEmpty()) {
+      return;
+    }
+
+    final List<OpcUaWriteRequest> missingNodeWriteRequests =
+        getMissingNodeWriteRequests(writeRequests, writeValuesOnce(writeRequests));
+    if (missingNodeWriteRequests.isEmpty()) {
       return;
     }
 
-    writeValue(value, timestamp, opcDataType, currentQuality, segments, null);
+    addMissingNodes(missingNodeWriteRequests);
+    validateRetriedWrites(missingNodeWriteRequests, writeValuesOnce(missingNodeWriteRequests));
   }
 
-  private void writeValue(
-      final Object value,
-      final long timestamp,
-      final NodeId opcDataType,
-      final StatusCode currentQuality,
-      final String[] segments,
-      final @Nullable String name)
-      throws Exception {
-    final NodeId nodeId =
-        new NodeId(
-            NAME_SPACE_INDEX,
-            Objects.nonNull(name)
-                ? String.join("/", segments) + "/" + name
-                : String.join("/", segments));
-    final Variant variant = new Variant(value);
-    final DataValue dataValue =
-        new DataValue(variant, currentQuality, new DateTime(timestamp), new DateTime());
-    StatusCode writeStatus = client.writeValue(nodeId, dataValue).get();
-
-    if (writeStatus.getValue() == StatusCodes.Bad_NodeIdUnknown) {
-      final AddNodesResponse addStatus =
-          client.addNodes(getNodesToAdd(segments, name, opcDataType, variant)).get();
-      for (final AddNodesResult result : addStatus.getResults()) {
-        if (!result.getStatusCode().equals(StatusCode.GOOD)
-            && !(result.getStatusCode().getValue() == StatusCodes.Bad_NodeIdExists)) {
-          throw new PipeException(
-              DataNodePipeMessages.FAILED_TO_CREATE_NODES_AFTER_TRANSFER_DATA
-                  + addStatus
-                  + getErrorString(segments, name, opcDataType, value, writeStatus));
+  private List<OpcUaWriteRequest> getMissingNodeWriteRequests(
+      final List<OpcUaWriteRequest> writeRequests, final List<StatusCode> writeStatuses) {
+    final List<OpcUaWriteRequest> missingNodeWriteRequests = new ArrayList<>();
+    for (int i = 0; i < writeRequests.size(); ++i) {
+      if (writeStatuses.get(i).getValue() == StatusCodes.Bad_NodeIdUnknown) {
+        missingNodeWriteRequests.add(writeRequests.get(i));
+      } else {
+        validateInitialWrite(writeRequests.get(i), writeStatuses.get(i));
+      }
+    }
+    return missingNodeWriteRequests;
+  }
+
+  private void validateInitialWrite(
+      final OpcUaWriteRequest writeRequest, final StatusCode writeStatus) {
+    if (writeStatus.getValue() != StatusCode.GOOD.getValue()) {
+      throw new PipeException(
+          DataNodePipeMessages.FAILED_TO_TRANSFER_DATAVALUE
+              + writeRequest.getErrorString(writeStatus));
+    }
+  }
+
+  private void addMissingNodes(final List<OpcUaWriteRequest> writeRequests) throws Exception {
+    final List<AddNodesItem> nodesToAdd = new ArrayList<>();
+    final Set<ExpandedNodeId> nodeIdsToAdd = new HashSet<>();
+    for (final OpcUaWriteRequest writeRequest : writeRequests) {
+      for (final AddNodesItem nodeToAdd :
+          getNodesToAdd(
+              writeRequest.segments,
+              writeRequest.name,
+              writeRequest.opcDataType,
+              writeRequest.variant)) {
+        if (nodeIdsToAdd.add(nodeToAdd.getRequestedNewNodeId())) {
+          nodesToAdd.add(nodeToAdd);
         }
       }
-      writeStatus = client.writeValue(nodeId, dataValue).get();
-      if (writeStatus.getValue() != StatusCode.GOOD.getValue()) {
+    }
+
+    final AddNodesResponse addStatus = client.addNodes(nodesToAdd).get();
+    for (final AddNodesResult result : addStatus.getResults()) {
+      if (!result.getStatusCode().equals(StatusCode.GOOD)
+          && result.getStatusCode().getValue() != StatusCodes.Bad_NodeIdExists) {
+        throw new PipeException(
+            DataNodePipeMessages.FAILED_TO_CREATE_NODES_AFTER_TRANSFER_DATA
+                + addStatus
+                + writeRequests
+                    .get(0)
+                    .getErrorString(new StatusCode(StatusCodes.Bad_NodeIdUnknown)));
+      }
+    }
+  }
+
+  private void validateRetriedWrites(
+      final List<OpcUaWriteRequest> writeRequests, final List<StatusCode> writeStatuses) {
+    for (int i = 0; i < writeRequests.size(); ++i) {
+      if (writeStatuses.get(i).getValue() != StatusCode.GOOD.getValue()) {
         throw new PipeException(
             DataNodePipeMessages.FAILED_TO_TRANSFER_DATAVALUE_AFTER_SUCCESSFULLY_CREATED
-                + getErrorString(segments, name, opcDataType, value, writeStatus));
+                + writeRequests.get(i).getErrorString(writeStatuses.get(i)));
       }
-    } else if (writeStatus.getValue() != StatusCode.GOOD.getValue()) {
-      throw new PipeException(
-          DataNodePipeMessages.FAILED_TO_TRANSFER_DATAVALUE
-              + getErrorString(segments, name, opcDataType, value, writeStatus));
+    }
+  }
+
+  private List<StatusCode> writeValuesOnce(final List<OpcUaWriteRequest> writeRequests)
+      throws Exception {
+    final List<NodeId> nodeIds = new ArrayList<>(writeRequests.size());
+    final List<DataValue> dataValues = new ArrayList<>(writeRequests.size());
+    for (final OpcUaWriteRequest writeRequest : writeRequests) {
+      nodeIds.add(writeRequest.nodeId);
+      dataValues.add(writeRequest.dataValue);
+    }
+    return client.writeValues(nodeIds, dataValues).get();
+  }
+
+  private static final class OpcUaWriteRequest {
+    private final Object value;
+    private final NodeId opcDataType;
+    private final String[] segments;
+    private final @Nullable String name;
+    private final NodeId nodeId;
+    private final Variant variant;
+    private final DataValue dataValue;
+
+    private OpcUaWriteRequest(
+        final Object value,
+        final long timestamp,
+        final NodeId opcDataType,
+        final StatusCode currentQuality,
+        final String[] segments,
+        final @Nullable String name) {
+      this.value = value;
+      this.opcDataType = opcDataType;
+      this.segments = segments;
+      this.name = name;
+      nodeId =
+          new NodeId(
+              NAME_SPACE_INDEX,
+              Objects.nonNull(name)
+                  ? String.join("/", segments) + "/" + name
+                  : String.join("/", segments));
+      variant = new Variant(value);
+      dataValue = new DataValue(variant, currentQuality, new DateTime(timestamp), new DateTime());
+    }
+
+    private String getErrorString(final StatusCode writeStatus) {
+      return IoTDBOpcUaClient.getErrorString(segments, name, opcDataType, value, writeStatus);
     }
   }
 

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/opcua/client/IoTDBOpcUaClientTest.java

@@ -0,0 +1,176 @@
+/*
+ * 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.iotdb.db.pipe.sink.protocol.opcua.client;
+
+import org.apache.iotdb.db.pipe.sink.protocol.opcua.OpcUaSink;
+import org.apache.iotdb.pipe.api.exception.PipeException;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.write.record.Tablet;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.eclipse.milo.opcua.sdk.client.OpcUaClient;
+import org.eclipse.milo.opcua.sdk.client.api.UaClient;
+import org.eclipse.milo.opcua.sdk.client.api.identity.AnonymousProvider;
+import org.eclipse.milo.opcua.stack.core.StatusCodes;
+import org.eclipse.milo.opcua.stack.core.security.SecurityPolicy;
+import org.eclipse.milo.opcua.stack.core.types.builtin.NodeId;
+import org.eclipse.milo.opcua.stack.core.types.builtin.StatusCode;
+import org.eclipse.milo.opcua.stack.core.types.structured.AddNodesItem;
+import org.eclipse.milo.opcua.stack.core.types.structured.AddNodesResponse;
+import org.eclipse.milo.opcua.stack.core.types.structured.AddNodesResult;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.ArgumentMatcher;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+public class IoTDBOpcUaClientTest {
+
+  @Test
+  public void testTransferWritesAllMeasurementsInOneRequest() throws Exception {
+    final OpcUaClient miloClient = Mockito.mock(OpcUaClient.class);
+    Mockito.when(miloClient.writeValues(Mockito.anyList(), Mockito.anyList()))
+        .thenReturn(
+            CompletableFuture.completedFuture(Arrays.asList(StatusCode.GOOD, StatusCode.GOOD)));
+    final IoTDBOpcUaClient client = createClient(miloClient);
+
+    client.transfer(createTablet(), createSink());
+
+    Mockito.verify(miloClient)
+        .writeValues(
+            Mockito.argThat(nodeIds("root/db/d1/s1", "root/db/d1/s2")),
+            Mockito.argThat(listWithSize(2)));
+  }
+
+  @Test
+  public void testTransferCreatesAndRetriesOnlyMissingNodes() throws Exception {
+    final OpcUaClient miloClient = Mockito.mock(OpcUaClient.class);
+    Mockito.when(miloClient.writeValues(Mockito.anyList(), Mockito.anyList()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Arrays.asList(new StatusCode(StatusCodes.Bad_NodeIdUnknown), StatusCode.GOOD)))
+        .thenReturn(CompletableFuture.completedFuture(Collections.singletonList(StatusCode.GOOD)));
+
+    final AddNodesResponse addNodesResponse = Mockito.mock(AddNodesResponse.class);
+    final AddNodesResult addNodesResult = Mockito.mock(AddNodesResult.class);
+    Mockito.when(addNodesResult.getStatusCode()).thenReturn(StatusCode.GOOD);
+    Mockito.when(addNodesResponse.getResults()).thenReturn(new AddNodesResult[] {addNodesResult});
+    Mockito.when(miloClient.addNodes(Mockito.anyList()))
+        .thenReturn(CompletableFuture.completedFuture(addNodesResponse));
+
+    final IoTDBOpcUaClient client = Mockito.spy(createClient(miloClient));
+    final AddNodesItem nodeToAdd = Mockito.mock(AddNodesItem.class);
+    Mockito.doReturn(Arrays.asList(nodeToAdd, nodeToAdd))
+        .when(client)
+        .getNodesToAdd(
+            Mockito.any(String[].class),
+            Mockito.eq("s1"),
+            Mockito.any(NodeId.class),
+            Mockito.any());
+
+    client.transfer(createTablet(), createSink());
+
+    final InOrder inOrder = Mockito.inOrder(miloClient);
+    inOrder
+        .verify(miloClient)
+        .writeValues(Mockito.argThat(listWithSize(2)), Mockito.argThat(listWithSize(2)));
+    inOrder.verify(miloClient).addNodes(Mockito.argThat(listWithSize(1)));
+    inOrder
+        .verify(miloClient)
+        .writeValues(Mockito.argThat(nodeIds("root/db/d1/s1")), Mockito.argThat(listWithSize(1)));
+  }
+
+  @Test
+  public void testTransferFailsOnNonRecoverableStatus() throws Exception {
+    final OpcUaClient miloClient = Mockito.mock(OpcUaClient.class);
+    Mockito.when(miloClient.writeValues(Mockito.anyList(), Mockito.anyList()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Arrays.asList(new StatusCode(StatusCodes.Bad_NotWritable), StatusCode.GOOD)));
+    final IoTDBOpcUaClient client = createClient(miloClient);
+
+    try {
+      client.transfer(createTablet(), createSink());
+      Assert.fail();
+    } catch (final PipeException e) {
+      Assert.assertTrue(e.getMessage().contains("root.db.d1.s1"));
+      Assert.assertTrue(e.getMessage().contains("Bad_NotWritable"));
+    }
+
+    Mockito.verify(miloClient, Mockito.never()).addNodes(Mockito.anyList());
+  }
+
+  private static IoTDBOpcUaClient createClient(final OpcUaClient miloClient) throws Exception {
+    final IoTDBOpcUaClient client =
+        new IoTDBOpcUaClient(
+            "opc.tcp://127.0.0.1:12686", SecurityPolicy.None, new AnonymousProvider(), false);
+    final ClientRunner runner = Mockito.mock(ClientRunner.class);
+    Mockito.when(runner.getTimeoutSeconds()).thenReturn(1L);
+    client.setRunner(runner);
+    final CompletableFuture<UaClient> connectFuture = CompletableFuture.completedFuture(miloClient);
+    Mockito.when(miloClient.connect()).thenReturn(connectFuture);
+    client.run(miloClient);
+    return client;
+  }
+
+  private static OpcUaSink createSink() {
+    final OpcUaSink sink = Mockito.mock(OpcUaSink.class);
+    Mockito.when(sink.getDefaultQuality()).thenReturn(StatusCode.GOOD);
+    return sink;
+  }
+
+  private static Tablet createTablet() {
+    final Tablet tablet =
+        new Tablet(
+            "root.db.d1",
+            Arrays.asList(
+                new MeasurementSchema("s1", TSDataType.INT64),
+                new MeasurementSchema("s2", TSDataType.DOUBLE)),
+            1);
+    tablet.addTimestamp(0, 1L);
+    tablet.addValue("s1", 0, 1L);
+    tablet.addValue("s2", 0, 2.0D);
+    tablet.setRowSize(1);
+    return tablet;
+  }
+
+  private static ArgumentMatcher<List<NodeId>> nodeIds(final String... identifiers) {
+    return nodeIds -> {
+      if (nodeIds.size() != identifiers.length) {
+        return false;
+      }
+      for (int i = 0; i < identifiers.length; ++i) {
+        if (!new NodeId(2, identifiers[i]).equals(nodeIds.get(i))) {
+          return false;
+        }
+      }
+      return true;
+    };
+  }
+
+  private static <T> ArgumentMatcher<List<T>> listWithSize(final int size) {
+    return values -> values.size() == size;
+  }
+}