com.huaweicloud.dws.client.collector.SecurityModeCollector Maven / Gradle / Ivy

Go to download
package com.huaweicloud.dws.client.collector;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.action.SecurityModeAction;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.exception.InvalidException;
import com.huaweicloud.dws.client.model.OperationType;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.model.WriteMode;
import com.huaweicloud.dws.client.worker.ExecutionPool;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @ProjectName: dws-connector
 * @ClassName: TableActionCollector
 * @Description: 表事件收集， 用于按表维处理事件
 * @Date: 2023/1/10 15:48
 * @Version: 1.0
 */
@Slf4j
public class SecurityModeCollector extends TableActionCollector {


    public SecurityModeCollector(DwsConfig config, ExecutionPool pool, TableSchema schema) {
        super(config, pool, schema);
    }

    private boolean hasCopy;

    private CopyReader copyReader;

    private SecurityModeAction currentAction;

    private List columnList;

    /**
     * 按批次维度缓存数据去重
     */
    private final Map nextRecord = new HashMap<>(30000);

    private boolean isWriteBuffer;


    @Override
    public void collector(Record record) throws DwsClientException {
        // 刷库后就会将字段清空，重新缓存
        if (columnList == null) {
            initEnv(record);
            // 是否需要写入缓存：如果不是纯写入，由于需要主键判断是否重复，必须写缓存，如果有重试也必须写缓存，因为需要获取全量数据
            isWriteBuffer = !(tableConfig.getWriteMode() == WriteMode.COPY || !record.isUpsert()) || config.getMaxFlushRetryTimes() != 1;
        }
        if (!hasCopy) {
            super.collector(record);
            return;
        }
        if (record.getType() != OperationType.WRITE) {
            // 流式写入不支持删除
            throw new InvalidException("security mode use copy write only support write record.");
        }
        pool.checkStatus();
        lock.lock();
        try {
            initCopy(record, null);
            if (isWriteBuffer) {
                // 如果当前数据存在过，放到下批入库
                Record write = buffer.write(record);
                if (write != null) {
                    nextRecord.put(write.getRecordKey(), write);
                    return;
                }
            }
            // 该模式下格式不会变化，可以直接将数据写入copy流
            copyReader.writeRecord(record);
            if (needFlush()) {
                // 只有在数量超过容量规定倍数后 才强制业务线程提交刷库，否则等待定时任务刷库即可
                flush(false);
            }
        } finally {
            lock.unlock();
        }

    }

    private boolean needFlush() {
        if (buffer.flush() && buffer.getSize() >= tableConfig.getAutoFlushBatchSize() * tableConfig.getBatchOutWeighRatio()) {
            return true;
        }
        return copyReader != null && copyReader.willStop() && copyReader.getSize() >= tableConfig.getAutoFlushBatchSize() * tableConfig.getBatchOutWeighRatio();
    }

    private synchronized void initEnv(Record record) {
        if (columnList != null) {
            return;
        }
        if (currentTableSchema == null) {
            currentTableSchema = record.getTableSchema();
            tableConfig = config.getTableConfig(currentTableSchema.getTableName());
        }
        hasCopy = tableConfig.getWriteMode().name().startsWith("COPY");
        columnList = new ArrayList<>();
        record.getColumnBit().stream().forEach(item -> columnList.add(currentTableSchema.getColumnNames().get(item)));
    }

    private void initCopy(Record record, Object lock) throws DwsClientException {
        if (copyReader != null) {
            return;
        }
        copyReader = new CopyReader(tableConfig);
        if (lock == null) {
            currentAction = new SecurityModeAction(record, config, copyReader, columnList);
        } else {
            currentAction = new SecurityModeAction(record, config, copyReader, columnList, lock);
        }
        long start = System.currentTimeMillis();
        while (!pool.submit(currentAction)) {
            log.debug("try submit.");
        }
        log.info("submit action successful. use time = {}", System.currentTimeMillis() - start);
    }

    @Override
    public void flush() throws DwsClientException {
        flush(true);
    }

    public void flush(boolean force) throws DwsClientException {
        if (!hasCopy) {
            super.flush();
            return;
        }
        lock.lock();
        try {
            if (copyReader == null || currentAction == null) {
                return;
            }
            // 结算完成上一个正在执行的copy
            log.info("will flush size {}", buffer.getSize());
            copyReader.setAllRecords(buffer.getRecords());
            copyReader.stop();
            // 清空当前copy
            copyReader = null;
            // 清空缓存
            buffer.clear();
            pool.waitAction(currentAction);
            if (!nextRecord.isEmpty()) {
                log.info("current batch has more record, size {}", nextRecord.size());
                for (Record record : nextRecord.values()) {
                    if (copyReader == null) {
                        initCopy(record, currentAction.getLock());
                    }
                    copyReader.writeRecord(record);
                    if (!force) {
                        buffer.write(record);
                    }
                }
                // 如果不是程序运行结束需要强制结算，则将重复数据直接放入下批中写入
                nextRecord.clear();
                if (force) {
                    copyReader.stop();
                    pool.waitAction(currentAction);
                    copyReader = null;
                }
            }
        } finally {
            lock.unlock();
        }
    }


    /**
     * 根据 缓存是否达到刷库要求决定是否需要刷库
     * @return
     */
    @Override
    public void tryFlush() throws DwsClientException {
        if (buffer.flush() || (copyReader != null && copyReader.willStop())) {
            flush(false);
        }
    }

    @Override
    public TableSchema getTableSchema() {
        return currentTableSchema;
    }
}