Make blacklist persistent

Signed-off-by: Evgeniy Zuykin <[email protected]>

fe/fe-core/src/main/java/com/starrocks/meta/BlackListSql.java

@@ -14,6 +14,7 @@
 
 package com.starrocks.meta;
 
+import java.util.Objects;
 import java.util.regex.Pattern;
 
 public class BlackListSql {
@@ -24,4 +25,27 @@ public BlackListSql(Pattern pattern, long id) {
 
     public Pattern pattern;
     public long id;
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        BlackListSql that = (BlackListSql) o;
+        if (id != that.id) {
+            return false;
+        }
+        if (this.pattern == null) {
+            return that.pattern == null;
+        }
+        return Objects.equals(pattern.pattern(), that.pattern.pattern());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(pattern, id);
+    }
 }

fe/fe-core/src/main/java/com/starrocks/meta/SqlBlackList.java

@@ -15,57 +15,124 @@
 
 package com.starrocks.meta;
 
+import com.staros.util.LockCloseable;
 import com.starrocks.common.AnalysisException;
 import com.starrocks.common.ErrorCode;
 import com.starrocks.common.ErrorReport;
+import com.starrocks.persist.AddSqlBlackList;
+import com.starrocks.persist.ImageWriter;
+import com.starrocks.persist.metablock.SRMetaBlockEOFException;
+import com.starrocks.persist.metablock.SRMetaBlockException;
+import com.starrocks.persist.metablock.SRMetaBlockID;
+import com.starrocks.persist.metablock.SRMetaBlockReader;
+import com.starrocks.persist.metablock.SRMetaBlockWriter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
+import java.io.IOException;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 // Used by sql's blacklist
 public class SqlBlackList {
-    private static final SqlBlackList INSTANCE = new SqlBlackList();
 
-    public static SqlBlackList getInstance() {
-        return INSTANCE;
-    }
+    private static final Logger LOG = LogManager.getLogger(SqlBlackList.class);
 
-    public static void verifying(String sql) throws AnalysisException {
+    public void verifying(String sql) throws AnalysisException {
         String formatSql = sql.replace("\r", " ").replace("\n", " ").replaceAll("\\s+", " ");
-        for (BlackListSql patternAndId : getInstance().sqlBlackListMap.values()) {
-            Matcher m = patternAndId.pattern.matcher(formatSql);
-            if (m.find()) {
-                ErrorReport.reportAnalysisException(ErrorCode.ERR_SQL_IN_BLACKLIST_ERROR);
+        try (LockCloseable ignored = new LockCloseable(rwLock.readLock())) {
+            for (BlackListSql patternAndId : sqlBlackListMap.values()) {
+                Matcher m = patternAndId.pattern.matcher(formatSql);
+                if (m.find()) {
+                    ErrorReport.reportAnalysisException(ErrorCode.ERR_SQL_IN_BLACKLIST_ERROR);
+                }
+            }
+        }
+    }
+
+    public void load(SRMetaBlockReader reader) throws IOException, SRMetaBlockException, SRMetaBlockEOFException {
+        try (LockCloseable ignored = new LockCloseable(rwLock.writeLock())) {
+            int cnt = reader.readInt();
+            for (int i = 0; i < cnt; i++) {
+                AddSqlBlackList addSqlBlackList = reader.readJson(AddSqlBlackList.class);
+                put(addSqlBlackList.id, Pattern.compile(addSqlBlackList.pattern));
             }
+            LOG.info("loaded {} SQL blacklist patterns", sqlBlackListMap.size());
         }
     }
 
     // we use string of sql as key, and (pattern, id) as value.
-    public void put(Pattern pattern) {
-        if (!sqlBlackListMap.containsKey(pattern.toString())) {
-            long id = ids.getAndIncrement();
-            sqlBlackListMap.putIfAbsent(pattern.toString(), new BlackListSql(pattern, id));
+    public long put(Pattern pattern) {
+        try (LockCloseable ignored = new LockCloseable(rwLock.writeLock())) {
+            BlackListSql blackListSql = sqlBlackListMap.get(pattern.toString());
+            if (blackListSql == null) {
+                long id = ids.getAndIncrement();
+                sqlBlackListMap.put(pattern.toString(), new BlackListSql(pattern, id));
+                return id;
+            } else {
+                return blackListSql.id;
+            }
+        }
+    }
+
+    public void put(long id, Pattern pattern) {
+        try (LockCloseable ignored = new LockCloseable(rwLock.writeLock())) {
+            BlackListSql blackListSql = sqlBlackListMap.get(pattern.toString());
+            if (blackListSql == null) {
+                ids.set(Math.max(ids.get(), id + 1));
+                sqlBlackListMap.put(pattern.toString(), new BlackListSql(pattern, id));
+            }
         }
     }
 
     // we delete sql's regular expression use id, so we iterate this map.
     public void delete(long id) {
-        for (Map.Entry<String, BlackListSql> entry : sqlBlackListMap.entrySet()) {
-            if (entry.getValue().id == id) {
-                sqlBlackListMap.remove(entry.getKey());
+        try (LockCloseable ignored = new LockCloseable(rwLock.writeLock())) {
+            for (Map.Entry<String, BlackListSql> entry : sqlBlackListMap.entrySet()) {
+                if (entry.getValue().id == id) {
+                    sqlBlackListMap.remove(entry.getKey());
+                }
             }
         }
     }
 
+    public void save(ImageWriter imageWriter) throws IOException, SRMetaBlockException {
+        try (LockCloseable ignored = new LockCloseable(rwLock.readLock())) {
+            // one for self and N for patterns
+            final int cnt = 1 + sqlBlackListMap.size();
+            SRMetaBlockWriter writer = imageWriter.getBlockWriter(SRMetaBlockID.BLACKLIST_MGR, cnt);
+
+            // write patterns
+            writer.writeInt(sqlBlackListMap.size());
+            for (BlackListSql p : sqlBlackListMap.values()) {
+                writer.writeJson(new AddSqlBlackList(p.id, p.pattern.pattern()));
+            }
+            writer.close();
+        }
+    }
+
+    public List<BlackListSql> getBlackLists() {
+        try (LockCloseable ignored = new LockCloseable(rwLock.readLock())) {
+            return this.sqlBlackListMap.values().stream().sorted(Comparator.comparing(x -> x.id)).collect(Collectors.toList());
+        }
+    }
+
+    private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
+
     // sqlBlackListMap: key is String(sql), value is BlackListSql.
     // BlackListSql is (Pattern, id). Pattern is the regular expression, id marks this sql, and is show with "show sqlblacklist";
-    public ConcurrentMap<String, BlackListSql> sqlBlackListMap = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, BlackListSql> sqlBlackListMap = new ConcurrentHashMap<>();
 
     // ids used in sql blacklist
-    public AtomicLong ids = new AtomicLong();
+    private final AtomicLong ids = new AtomicLong();
 }
 

fe/fe-core/src/main/java/com/starrocks/persist/AddSqlBlackList.java

@@ -0,0 +1,50 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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 com.starrocks.persist;
+
+import com.google.gson.annotations.SerializedName;
+import com.starrocks.common.io.JsonWriter;
+
+import java.util.Objects;
+
+public class AddSqlBlackList extends JsonWriter {
+    public AddSqlBlackList(long id, String pattern) {
+        this.id = id;
+        this.pattern = pattern;
+    }
+
+    @SerializedName("id")
+    public final long id;
+
+    @SerializedName("pattern")
+    public final String pattern;
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        AddSqlBlackList that = (AddSqlBlackList) o;
+        return id == that.id && Objects.equals(pattern, that.pattern);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id, pattern);
+    }
+}

fe/fe-core/src/main/java/com/starrocks/persist/DeleteSqlBlackLists.java

@@ -0,0 +1,47 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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 com.starrocks.persist;
+
+import com.google.gson.annotations.SerializedName;
+import com.starrocks.common.io.JsonWriter;
+
+import java.util.List;
+import java.util.Objects;
+
+public class DeleteSqlBlackLists extends JsonWriter {
+    public DeleteSqlBlackLists(List<Long> ids) {
+        this.ids = ids;
+    }
+
+    @SerializedName("ids")
+    public final List<Long> ids;
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        DeleteSqlBlackLists that = (DeleteSqlBlackLists) o;
+        return Objects.equals(ids, that.ids);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(ids);
+    }
+}

fe/fe-core/src/main/java/com/starrocks/persist/EditLog.java

@@ -124,6 +124,7 @@
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutionException;
+import java.util.regex.Pattern;
 
 /**
  * EditLog maintains a log of the memory modifications.
@@ -1111,6 +1112,19 @@ public void loadJournal(GlobalStateMgr globalStateMgr, JournalEntity journal)
                     globalStateMgr.getClusterSnapshotMgr().replayLog(log);
                     break;
                 }
+                case OperationType.OP_ADD_SQL_QUERY_BLACK_LIST: {
+                    AddSqlBlackList addBlacklistRequest = (AddSqlBlackList) journal.data();
+                    GlobalStateMgr.getCurrentState().getSqlBlackList()
+                            .put(addBlacklistRequest.id, Pattern.compile(addBlacklistRequest.pattern));
+                    break;
+                }
+                case OperationType.OP_DELETE_SQL_QUERY_BLACK_LIST: {
+                    DeleteSqlBlackLists deleteBlackListsRequest = (DeleteSqlBlackLists) journal.data();
+                    for (int i = 0; i < deleteBlackListsRequest.ids.size(); i++) {
+                        GlobalStateMgr.getCurrentState().getSqlBlackList().delete(deleteBlackListsRequest.ids.get(i));
+                    }
+                    break;
+                }
                 default: {
                     if (Config.metadata_ignore_unknown_operation_type) {
                         LOG.warn("UNKNOWN Operation Type {}", opCode);
@@ -1795,6 +1809,15 @@ public void logAlterMaterializedViewProperties(ModifyTablePropertyOperationLog l
         logEdit(OperationType.OP_ALTER_MATERIALIZED_VIEW_PROPERTIES, log);
     }
 
+    public void logAddSQLBlackList(AddSqlBlackList addBlackList) {
+        logEdit(OperationType.OP_ADD_SQL_QUERY_BLACK_LIST, addBlackList);
+    }
+
+    public void logDeleteSQLBlackList(DeleteSqlBlackLists deleteBlacklists) {
+        logEdit(OperationType.OP_DELETE_SQL_QUERY_BLACK_LIST, deleteBlacklists);
+    }
+
+
     public void logStarMgrOperation(StarMgrJournal journal) {
         logEdit(OperationType.OP_STARMGR, journal);
     }

fe/fe-core/src/main/java/com/starrocks/persist/EditLogDeserializer.java

@@ -243,6 +243,8 @@ public class EditLogDeserializer {
             .put(OperationType.OP_ALTER_WAREHOUSE, Warehouse.class)
             .put(OperationType.OP_DROP_WAREHOUSE, DropWarehouseLog.class)
             .put(OperationType.OP_CLUSTER_SNAPSHOT_LOG, ClusterSnapshotLog.class)
+            .put(OperationType.OP_ADD_SQL_QUERY_BLACK_LIST, AddSqlBlackList.class)
+            .put(OperationType.OP_DELETE_SQL_QUERY_BLACK_LIST, DeleteSqlBlackLists.class)
             .build();
 
     public static Writable deserialize(Short opCode, DataInput in) throws IOException {

fe/fe-core/src/main/java/com/starrocks/persist/OperationType.java

@@ -591,6 +591,11 @@ public class OperationType {
     @IgnorableOnReplayFailed
     public static final short OP_CLUSTER_SNAPSHOT_LOG = 13513;
 
+    @IgnorableOnReplayFailed
+    public static final short OP_ADD_SQL_QUERY_BLACK_LIST = 13700;
+    @IgnorableOnReplayFailed
+    public static final short OP_DELETE_SQL_QUERY_BLACK_LIST = 13701;
+
     /**
      * NOTICE: OperationType cannot use a value exceeding 20000, please follow the above sequence number
      */

fe/fe-core/src/main/java/com/starrocks/persist/metablock/SRMetaBlockID.java

@@ -99,6 +99,8 @@ public int getId() {
 
     public static final SRMetaBlockID CLUSTER_SNAPSHOT_MGR = new SRMetaBlockID(33);
 
+    public static final SRMetaBlockID BLACKLIST_MGR = new SRMetaBlockID(34);
+
     /**
      * NOTICE: SRMetaBlockID cannot use a value exceeding 20000, please follow the above sequence number
      */

fe/fe-core/src/main/java/com/starrocks/qe/ShowExecutor.java

@@ -129,7 +129,6 @@
 import com.starrocks.load.streamload.StreamLoadFunctionalExprProvider;
 import com.starrocks.load.streamload.StreamLoadTask;
 import com.starrocks.meta.BlackListSql;
-import com.starrocks.meta.SqlBlackList;
 import com.starrocks.proto.FailPointTriggerModeType;
 import com.starrocks.proto.PFailPointInfo;
 import com.starrocks.proto.PFailPointTriggerMode;
@@ -2234,10 +2233,10 @@ public ShowResultSet visitShowPluginsStatement(ShowPluginsStmt statement, Connec
         @Override
         public ShowResultSet visitShowSqlBlackListStatement(ShowSqlBlackListStmt statement, ConnectContext context) {
             List<List<String>> rows = new ArrayList<>();
-            for (Map.Entry<String, BlackListSql> entry : SqlBlackList.getInstance().sqlBlackListMap.entrySet()) {
+            for (BlackListSql entry : GlobalStateMgr.getCurrentState().getSqlBlackList().getBlackLists()) {
                 List<String> oneSql = new ArrayList<>();
-                oneSql.add(String.valueOf(entry.getValue().id));
-                oneSql.add(entry.getKey());
+                oneSql.add(String.valueOf(entry.id));
+                oneSql.add(entry.pattern.toString());
                 rows.add(oneSql);
             }
             return new ShowResultSet(statement.getMetaData(), rows);