com.alibaba.hologres.client.RecordInputFormat Maven / Gradle / Ivy
/*
* Copyright (c) 2021. Alibaba Group Holding Limited
*/
package com.alibaba.hologres.client;
import com.alibaba.hologres.client.exception.HoloClientException;
import com.alibaba.hologres.client.impl.RecordReader;
import com.alibaba.hologres.client.model.ExportContext;
import com.alibaba.hologres.client.model.Record;
import com.alibaba.hologres.client.model.TableSchema;
import org.postgresql.jdbc.TimestampUtils;
import java.io.IOException;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
/**
* RecordInputStream.
*/
public class RecordInputFormat {
private final TableSchema schema;
private final TimestampUtils timestampUtils;
private final ExportContext exportContext;
public RecordInputFormat(ExportContext exportContext, TableSchema schema) throws IOException {
this.exportContext = exportContext;
this.schema = schema;
this.threadSize = exportContext.getThreadSize();
this.numOpened = new AtomicInteger(threadSize);
this.timestampUtils = exportContext.getTimestampUtils();
this.queue = new ArrayBlockingQueue(1024);
this.start();
}
int threadSize;
AtomicInteger numOpened;
BlockingQueue queue;
ExecutorService threadPool = Executors.newCachedThreadPool();
//call getRecord until null or call RecordInputFormat.cancel() to interrupt
public Record getRecord() {
while (numOpened.get() > 0 || !queue.isEmpty()) {
Record r;
try {
r = queue.poll(100, TimeUnit.MILLISECONDS);
} catch (InterruptedException e) {
return null;
}
if (r != null) {
return r;
}
}
threadPool.shutdown();
return null;
}
private void start() {
for (int i = 0; i < threadSize; i++) {
threadPool.execute(new RecordReader(exportContext.getInputStream(i), schema, queue, numOpened, timestampUtils));
}
}
public void cancel() throws HoloClientException {
threadPool.shutdownNow();
numOpened.set(0);
exportContext.cancel();
}
}