@Override
public void open(Configuration parameters) throws Exception {
super.open(parameters);
KinesisProducerConfiguration producerConfig = new KinesisProducerConfiguration();
producerConfig.setRegion(configProps.getProperty(ProducerConfigConstants.AWS_REGION));
producerConfig.setCredentialsProvider(AWSUtil.getCredentialsProvider(configProps));
if (configProps.containsKey(ProducerConfigConstants.COLLECTION_MAX_COUNT)) {
producerConfig.setCollectionMaxCount(
PropertiesUtil.getLong(
configProps,
ProducerConfigConstants.COLLECTION_MAX_COUNT,
producerConfig.getCollectionMaxCount(),
LOG));
}
if (configProps.containsKey(ProducerConfigConstants.AGGREGATION_MAX_COUNT)) {
producerConfig.setAggregationMaxCount(
PropertiesUtil.getLong(
configProps,
ProducerConfigConstants.AGGREGATION_MAX_COUNT,
producerConfig.getAggregationMaxCount(),
LOG));
}
producer = new KinesisProducer(producerConfig);
callback =
new FutureCallback<UserRecordResult>() {
@Override
public void onSuccess(UserRecordResult result) {
if (!result.isSuccessful()) {
if (failOnError) {
thrownException = new RuntimeException("Record was not sent successful");
} else {
LOG.warn("Record was not sent successful");
}
}
}
@Override
public void onFailure(Throwable t) {
if (failOnError) {
thrownException = t;
} else {
LOG.warn("An exception occurred while processing a record", t);
}
}
};
if (this.customPartitioner != null) {
this.customPartitioner.initialize(
getRuntimeContext().getIndexOfThisSubtask(),
getRuntimeContext().getNumberOfParallelSubtasks());
}
LOG.info("Started Kinesis producer instance for region '{}'", producerConfig.getRegion());
}
@Override
public void close() throws Exception {
LOG.info("Closing producer");
super.close();
KinesisProducer kp = this.producer;
this.producer = null;
if (kp != null) {
LOG.info("Flushing outstanding {} records", kp.getOutstandingRecordsCount());
// try to flush all outstanding records
while (kp.getOutstandingRecordsCount() > 0) {
kp.flush();
try {
Thread.sleep(500);
} catch (InterruptedException e) {
LOG.warn("Flushing was interrupted.");
// stop the blocking flushing and destroy producer immediately
break;
}
}
LOG.info("Flushing done. Destroying producer instance.");
kp.destroy();
}
}