Skip to content

Commit 8e796f4

Browse files
HTHouHTHou
committed
[Feature] Exclude audit database from IoTConsensusV2 audit
1 parent fc06291 commit 8e796f4

4 files changed

Lines changed: 82 additions & 3 deletions

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/audit/DataNodeUserDataTransferAuditor.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,15 @@ public static boolean isEnabled() {
4646
return COMMON_CONFIG.isEnableAuditLog();
4747
}
4848

49+
public static boolean isEnabledFor(ConsensusGroupId consensusGroupId) {
50+
try {
51+
return isEnabled() && containsUserData(consensusGroupId);
52+
} catch (RuntimeException ignored) {
53+
// Classification is advisory and must not affect user data transfer.
54+
return false;
55+
}
56+
}
57+
4958
public static void record(
5059
TEndPoint initiator,
5160
TEndPoint source,

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/IoTConsensusV2AsyncSink.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.iotdb.commons.client.IClientManager;
2626
import org.apache.iotdb.commons.client.async.AsyncIoTConsensusV2ServiceClient;
2727
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
28+
import org.apache.iotdb.commons.consensus.DataRegionId;
2829
import org.apache.iotdb.commons.consensus.index.ProgressIndex;
2930
import org.apache.iotdb.commons.consensus.iotv2.container.IoTV2GlobalComponentContainer;
3031
import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkRetryTimesConfigurableException;
@@ -112,6 +113,7 @@ public class IoTConsensusV2AsyncSink extends IoTDBSink implements ConsensusPipeS
112113
private final int thisDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
113114
private final TEndPoint localEndPoint =
114115
new TEndPoint(IOTDB_CONFIG.getInternalAddress(), IOTDB_CONFIG.getDataRegionConsensusPort());
116+
private DataRegionId dataRegionId;
115117
private IoTConsensusV2SinkMetrics iotConsensusV2SinkMetrics;
116118
private String consensusPipeName;
117119
private int consensusGroupId;
@@ -141,6 +143,7 @@ public void customize(PipeParameters parameters, PipeConnectorRuntimeConfigurati
141143

142144
// Get consensusGroupId from parameters passed by IoTConsensusV2Impl
143145
consensusGroupId = parameters.getInt(CONNECTOR_CONSENSUS_GROUP_ID_KEY);
146+
dataRegionId = new DataRegionId(consensusGroupId);
144147
// Get consensusPipeName from parameters passed by IoTConsensusV2Impl
145148
consensusPipeName = parameters.getString(CONNECTOR_CONSENSUS_PIPE_NAME);
146149

@@ -750,7 +753,7 @@ public TEndPoint getFollowerUrl() {
750753
}
751754

752755
public boolean isUserDataTransferAuditEnabled() {
753-
return DataNodeUserDataTransferAuditor.isEnabled();
756+
return dataRegionId != null && DataNodeUserDataTransferAuditor.isEnabledFor(dataRegionId);
754757
}
755758

756759
public void recordUserDataTransferAudit(boolean success, String errorCode, Throwable error) {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/IoTConsensusV2SyncSink.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2626
import org.apache.iotdb.commons.client.IClientManager;
2727
import org.apache.iotdb.commons.client.sync.SyncIoTConsensusV2ServiceClient;
28+
import org.apache.iotdb.commons.consensus.DataRegionId;
2829
import org.apache.iotdb.commons.consensus.index.ProgressIndex;
2930
import org.apache.iotdb.commons.consensus.iotv2.container.IoTV2GlobalComponentContainer;
3031
import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkRetryTimesConfigurableException;
@@ -90,6 +91,7 @@ public class IoTConsensusV2SyncSink extends IoTDBSink {
9091
private final int consensusGroupId;
9192
private final IoTConsensusV2SinkMetrics iotConsensusV2SinkMetrics;
9293
private final TEndPoint localEndPoint;
94+
private final DataRegionId dataRegionId;
9395
private IoTConsensusV2SyncBatchReqBuilder tabletBatchBuilder;
9496

9597
public IoTConsensusV2SyncSink(
@@ -102,6 +104,7 @@ public IoTConsensusV2SyncSink(
102104
// retain the implementation of list to cope with possible future expansion
103105
this.peers = peers;
104106
this.consensusGroupId = consensusGroupId;
107+
this.dataRegionId = new DataRegionId(consensusGroupId);
105108
this.thisDataNodeId = thisDataNodeId;
106109
this.localEndPoint =
107110
new TEndPoint(
@@ -218,7 +221,7 @@ private void doTransfer() {
218221
resp.getBatchResps().stream()
219222
.map(TIoTConsensusV2TransferResp::getStatus)
220223
.collect(Collectors.toList());
221-
if (DataNodeUserDataTransferAuditor.isEnabled()) {
224+
if (isUserDataTransferAuditEnabled()) {
222225
final TSStatus failedStatus =
223226
statusList.stream().filter(status -> !isSuccessful(status)).findFirst().orElse(null);
224227
recordTransferAttempt(
@@ -582,13 +585,17 @@ private static boolean isSuccessful(TSStatus status) {
582585
}
583586

584587
private void recordTransferAttempt(boolean success, String errorCode, Throwable error) {
585-
if (!DataNodeUserDataTransferAuditor.isEnabled()) {
588+
if (!isUserDataTransferAuditEnabled()) {
586589
return;
587590
}
588591
DataNodeUserDataTransferAuditor.record(
589592
localEndPoint, localEndPoint, getFollowerUrl(), success, errorCode, error);
590593
}
591594

595+
private boolean isUserDataTransferAuditEnabled() {
596+
return DataNodeUserDataTransferAuditor.isEnabledFor(dataRegionId);
597+
}
598+
592599
// synchronized to avoid close connector when transfer event
593600
@Override
594601
public synchronized void close() {

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/audit/DataNodeUserDataTransferAuditorTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,24 @@
1919

2020
package org.apache.iotdb.db.audit;
2121

22+
import org.apache.iotdb.commons.conf.CommonConfig;
23+
import org.apache.iotdb.commons.conf.CommonDescriptor;
24+
import org.apache.iotdb.commons.consensus.DataRegionId;
2225
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
2326
import org.apache.iotdb.commons.request.IConsensusRequest;
2427
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
2528
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ObjectNode;
29+
import org.apache.iotdb.db.storageengine.StorageEngine;
30+
import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
2631

32+
import org.junit.After;
33+
import org.junit.Before;
2734
import org.junit.Test;
35+
import org.junit.runner.RunWith;
36+
import org.powermock.api.mockito.PowerMockito;
37+
import org.powermock.core.classloader.annotations.PowerMockIgnore;
38+
import org.powermock.core.classloader.annotations.PrepareForTest;
39+
import org.powermock.modules.junit4.PowerMockRunner;
2840

2941
import java.util.Collections;
3042

@@ -35,15 +47,63 @@
3547
import static org.mockito.Mockito.verify;
3648
import static org.mockito.Mockito.when;
3749

50+
@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
51+
@RunWith(PowerMockRunner.class)
52+
@PrepareForTest(StorageEngine.class)
3853
public class DataNodeUserDataTransferAuditorTest {
3954

55+
private final CommonConfig commonConfig = CommonDescriptor.getInstance().getConfig();
56+
private boolean auditLogEnabled;
57+
58+
@Before
59+
public void setUp() {
60+
auditLogEnabled = commonConfig.isEnableAuditLog();
61+
}
62+
63+
@After
64+
public void tearDown() {
65+
commonConfig.setEnableAuditLog(auditLogEnabled);
66+
}
67+
4068
@Test
4169
public void testAuditDatabaseIsExcludedFromGroupTransferAudit() {
4270
assertFalse(DataNodeUserDataTransferAuditor.containsUserData("__audit"));
4371
assertFalse(DataNodeUserDataTransferAuditor.containsUserData("root.__audit"));
4472
assertTrue(DataNodeUserDataTransferAuditor.containsUserData("root.sg"));
4573
}
4674

75+
@Test
76+
public void testIoTConsensusV2AuditGateExcludesAuditDatabase() {
77+
final StorageEngine storageEngine = mock(StorageEngine.class);
78+
final DataRegion auditDataRegion = mock(DataRegion.class);
79+
final DataRegion userDataRegion = mock(DataRegion.class);
80+
final DataRegionId auditDataRegionId = new DataRegionId(1);
81+
final DataRegionId userDataRegionId = new DataRegionId(2);
82+
PowerMockito.mockStatic(StorageEngine.class);
83+
PowerMockito.when(StorageEngine.getInstance()).thenReturn(storageEngine);
84+
when(storageEngine.getDataRegion(auditDataRegionId)).thenReturn(auditDataRegion);
85+
when(storageEngine.getDataRegion(userDataRegionId)).thenReturn(userDataRegion);
86+
when(auditDataRegion.getDatabaseName()).thenReturn("root.__audit");
87+
when(userDataRegion.getDatabaseName()).thenReturn("root.sg");
88+
89+
commonConfig.setEnableAuditLog(true);
90+
91+
assertFalse(DataNodeUserDataTransferAuditor.isEnabledFor(auditDataRegionId));
92+
assertTrue(DataNodeUserDataTransferAuditor.isEnabledFor(userDataRegionId));
93+
}
94+
95+
@Test
96+
public void testIoTConsensusV2AuditGateShortCircuitsWhenDisabled() {
97+
final StorageEngine storageEngine = mock(StorageEngine.class);
98+
PowerMockito.mockStatic(StorageEngine.class);
99+
PowerMockito.when(StorageEngine.getInstance()).thenReturn(storageEngine);
100+
commonConfig.setEnableAuditLog(false);
101+
102+
assertFalse(DataNodeUserDataTransferAuditor.isEnabledFor(new DataRegionId(1)));
103+
104+
verify(storageEngine, never()).getDataRegion(new DataRegionId(1));
105+
}
106+
47107
@Test
48108
public void testAuditDatabaseIsExcludedFromConsensusTransferAudit() {
49109
final InsertNode insertNode = mock(InsertNode.class);

0 commit comments

Comments
 (0)