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

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

import com.huaweicloud.dws.client.TableConfig;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.util.JdbcUtil;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.Reader;
import java.util.List;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @ProjectName: dws-connector
 * @ClassName: CopyManager
 * @Description:
 * @Date: 2023/12/8 16:32
 * @Version: 1.0
 */
@Slf4j
public class CopyReader extends Reader {
    private final LinkedBlockingQueue records = new LinkedBlockingQueue<>();

    private List allRecords;
    private String str;
    private int length;
    private int next = 0;
    private final TableConfig tableConfig;
    private final AtomicBoolean stop = new AtomicBoolean();
    private volatile int size;
    private volatile boolean reTry = false;

    private final StringBuilder temp = new StringBuilder();
    private final long startTime;

    public void setAllRecords(List allRecords) {
        this.allRecords = allRecords;
    }

    public void writeRecord(Record record) {
        while (!records.offer(record)) {
            log.debug("try submit.");
        }
        size++;
    }

    public CopyReader(TableConfig config) {
        this.tableConfig = config;
        startTime = System.currentTimeMillis();
    }

    @Override
    public int read(char[] cbuf, int off, int len) throws IOException {
        synchronized (lock) {
            boolean idxLess0 = off < 0 || len < 0;
            boolean idxMoreLen = off + len > cbuf.length;
            if (idxLess0 || idxMoreLen) {
                throw new IndexOutOfBoundsException();
            } else if (len == 0) {
                return 0;
            }
            if (reTry) {
                // 重试的时候直接拼一个大buffer
                str = JdbcUtil.buildCopyBuffer(allRecords, null, tableConfig);
                length = str.length();
            }
            if (next >= length) {
                if (stop.get() && records.isEmpty()) {
                    return -1;
                }
                if (records.size() > 1) {
                    int batch = Math.min(records.size(), 1000);
                    for (int i = 0; i < batch; i++) {
                        JdbcUtil.buildCopyBuffer(records.poll(), tableConfig, temp);
                    }
                } else {
                    try {
                        Record poll = records.poll(2, TimeUnit.MILLISECONDS);
                        if (poll == null) {
                            return 0;
                        }
                        JdbcUtil.buildCopyBuffer(poll, tableConfig, temp);
                    } catch (Exception e) {
                        log.error("", e);
                    }
                }
                str = temp.toString();
                temp.setLength(0);
                length = str.length();
            }
            return readNext(cbuf, off, len);
        }
    }

    private int readNext(char[] cbuf, int off, int len) {
        int leftLen = length - next;
        int n = Math.min(leftLen, len);
        str.getChars(next, next + n, cbuf, off);
        if (leftLen > len) {
            next += n;
        } else {
            next = 0;
            length = 0;
        }
        return n;
    }

    @Override
    public void close() throws IOException {
    }

    public void stop() {
        log.info("will stop size is {}.", size);
        stop.set(true);
    }

    public boolean willStop() {
        if (tableConfig.getAutoFlushMaxIntervalMs() < System.currentTimeMillis() - startTime) {
            return true;
        }
        return tableConfig.getAutoFlushBatchSize() <= size;
    }

    public int getSize() {
        return size;
    }

    public void reTry() {
        this.reTry = true;
    }
}