1

public class QuickwitSinkTask extends BkMetricSink<String, String> {

    private static final Logger log = LoggerFactory.getLogger(QuickwitSinkTask.class);

    private final QuickwitSinkConfig config;
    private List<BkSinkRecord<String>> recordList;
    private List<String> valueList;
    private String url;
    private long batchSizeInMb;
    private long totalSize;
    private int batchSize;
    private int retryBackoffMs;
    private int maxRetries;
    private ScheduledExecutorService flushExecutor;

    public QuickwitSinkTask(QuickwitSinkConfig config) {
        this.config = config;
    }

    @Override
    protected void doOpen(BkDatabusContext context) {
        recordList = Lists.newArrayList();
        valueList = Lists.newArrayList();

        url = config.getUrl();
        batchSizeInMb = config.getBatchSizeInMb() << 20;
        batchSize = config.getBatchSize();
        retryBackoffMs = config.getRetryBackoffMs();
        maxRetries = config.getMaxRetries();
        long batchTimeMs = config.getBulkFlushIntervalInMs();
        totalSize = 0L;
        flushExecutor = Executors.newScheduledThreadPool(config.getTaskNum());
    }

    @Override
    protected DataBusResult<String> doWrite(BkSinkRecord<String> record) throws Exception {
        List<String> values = JsonUtils.getValue(record.getValue());
        int currentNum;
        synchronized (this) {
            valueList.addAll(values);
            recordList.add(record);
            currentNum = valueList.size();
            totalSize += values.stream().mapToLong(value -> value.getBytes(StandardCharsets.UTF_8).length).sum();
            if (currentNum >= batchSize || totalSize >= batchSizeInMb) {
                final List<String> value = valueList;
                final List<BkSinkRecord<String>> records = recordList;
                flushExecutor.execute(() -> sendData(value, records));
                valueList = Lists.newArrayList();
                recordList = Lists.newArrayList();
                totalSize = 0L;
            }
        }
        return null;
    }

    @Override
    protected String getClusterType() {
        return "quickwit";
    }

    @Override
    protected void doClose() {
        if (null != flushExecutor) {
            flushExecutor.shutdown();
        }
    }


    private void sendData(List<String> values, List<BkSinkRecord<String>> bkRecordList) {

        String records = String.join("\n", values);

        int retries = 0;
        while (retries < maxRetries) {
            try {
                HttpURLConnection connection = HttpUtils.post(url, records, null);
                int responseCode = connection.getResponseCode();
                if (responseCode >= 200 && responseCode < 300) {
                    bkRecordList.forEach(BkSourceRecord::ack);
                    break;
                } else {
                    throw new IOException("Failed to send data, response code: [" + responseCode + "], "
                            + "content : [" + connection.getErrorStream() + "]");
                }
            } catch (IOException e) {
                retries++;
                if (retries < maxRetries) {
                    try {
                        Thread.sleep(retryBackoffMs);
                    } catch (InterruptedException ex) {
                        bkRecordList.forEach(BkSourceRecord::fail);
                        throw new RuntimeException(ex);
                    }
                } else {
                    bkRecordList.forEach(BkSourceRecord::fail);
                    log.error("Failed to send data after max retries");
                }
            }
        }

        values = null;
        bkRecordList = null;
    }
}