@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(); } }