com.amazonaws.services.kinesis.producer.KinesisProducer Java Examples

private KinesisProducer createKinesisProducer(KinesisPublisherConfiguration configuration) {
    ProducerConfiguration pc = configuration.getProducer();
    KinesisProducerConfiguration kc = new KinesisProducerConfiguration();

    pc.isAggregationEnabled().ifPresent(kc::setAggregationEnabled);
    pc.getAggregationMaxCount().ifPresent(kc::setAggregationMaxCount);
    pc.getAggregationMaxSize().ifPresent(kc::setAggregationMaxSize);
    pc.getCloudwatchEndpoint().ifPresent(kc::setCloudwatchEndpoint);
    pc.getCloudwatchPort().ifPresent(kc::setCloudwatchPort);
    pc.getCollectionMaxCount().ifPresent(kc::setCollectionMaxCount);
    pc.getCollectionMaxSize().ifPresent(kc::setCollectionMaxSize);
    pc.getConnectTimeout().ifPresent(kc::setConnectTimeout);
    pc.isEnableCoreDumps().ifPresent(kc::setEnableCoreDumps);
    pc.isFailIfThrottled().ifPresent(kc::setFailIfThrottled);
    pc.getKinesisEndpoint().ifPresent(kc::setKinesisEndpoint);
    pc.getKinesisPort().ifPresent(kc::setKinesisPort);
    pc.getLogLevel().ifPresent(kc::setLogLevel);
    pc.getMaxConnections().ifPresent(kc::setMaxConnections);
    pc.getMinConnections().ifPresent(kc::setMinConnections);
    pc.getNativeExecutable().ifPresent(kc::setNativeExecutable);
    pc.getRateLimit().ifPresent(kc::setRateLimit);
    pc.getRecordMaxBufferedTime().ifPresent(kc::setRecordMaxBufferedTime);
    pc.getRecordTtl().ifPresent(kc::setRecordTtl);
    pc.getRegion().ifPresent(kc::setRegion);
    pc.getRequestTimeout().ifPresent(kc::setRequestTimeout);
    pc.getTempDirectory().ifPresent(kc::setTempDirectory);
    pc.isVerifyCertificate().ifPresent(kc::setVerifyCertificate);

    try {
        return new KinesisProducer(kc);
    } catch (Exception e) {
        LOG.error("Failed to initialize " + KinesisPublisher.class.getSimpleName(), e);
        return null;
    }
}
 

private AbstractAwsMessageHandler<?> createKplMessageHandler(ProducerDestination destination,
		FunctionExpression<Message<?>> partitionKeyExpression) {

	final KplMessageHandler messageHandler;
	messageHandler = new KplMessageHandler(new KinesisProducer(this.kinesisProducerConfiguration));
	messageHandler.setStream(destination.getName());
	messageHandler.setPartitionKeyExpression(partitionKeyExpression);
	return messageHandler;
}
 

private ListenableFuture<UserRecordResult> addUserRecord(KinesisProducer kp, String streamName, String partitionKey,
		boolean usePartitionAsHashKey, SinkRecord sinkRecord) {

	// If configured use kafka partition key as explicit hash key
	// This will be useful when sending data from same partition into
	// same shard
	if (usePartitionAsHashKey)
		return kp.addUserRecord(streamName, partitionKey, Integer.toString(sinkRecord.kafkaPartition()),
				DataUtility.parseValue(sinkRecord.valueSchema(), sinkRecord.value()));
	else
		return kp.addUserRecord(streamName, partitionKey,
				DataUtility.parseValue(sinkRecord.valueSchema(), sinkRecord.value()));

}
 

@Override
public void stop() {
	// destroying kinesis producers which were not closed as part of close
	if (singleKinesisProducerPerPartition) {
		for (KinesisProducer kp : producerMap.values()) {
			kp.flushSync();
			kp.destroy();
		}
	} else {
		kinesisProducer.destroy();
	}

}
 

private KinesisProducer getKinesisProducer() {
	KinesisProducerConfiguration config = new KinesisProducerConfiguration();
	config.setRegion(regionName);
	config.setCredentialsProvider(new DefaultAWSCredentialsProviderChain());
	config.setMaxConnections(maxConnections);

	config.setAggregationEnabled(aggregration);

	// Limits the maximum allowed put rate for a shard, as a percentage of
	// the
	// backend limits.
	config.setRateLimit(rateLimit);

	// Maximum amount of time (milliseconds) a record may spend being
	// buffered
	// before it gets sent. Records may be sent sooner than this depending
	// on the
	// other buffering limits
	config.setRecordMaxBufferedTime(maxBufferedTime);

	// Set a time-to-live on records (milliseconds). Records that do not get
	// successfully put within the limit are failed.
	config.setRecordTtl(ttl);

	// Controls the number of metrics that are uploaded to CloudWatch.
	// Expected pattern: none|summary|detailed
	config.setMetricsLevel(metricsLevel);

	// Controls the granularity of metrics that are uploaded to CloudWatch.
	// Greater granularity produces more metrics.
	// Expected pattern: global|stream|shard
	config.setMetricsGranularity(metricsGranuality);

	// The namespace to upload metrics under.
	config.setMetricsNamespace(metricsNameSpace);

	return new KinesisProducer(config);

}
 

/**
 * Here'll walk through some of the config options and create an instance of
 * KinesisProducer, which will be used to put records.
 *
 * @param region The region of the Kinesis stream being used.
 *
 * @return KinesisProducer instance used to put records.
 */
private static KinesisProducer getKinesisProducer(final String region) {
    KinesisProducerConfiguration config = new KinesisProducerConfiguration();
    config.setRegion(region);
    config.setCredentialsProvider(new DefaultAWSCredentialsProviderChain());
    config.setMaxConnections(1);
    config.setRequestTimeout(60000);
    config.setRecordMaxBufferedTime(2000);
    config.setAggregationEnabled(false);

    KinesisProducer producer = new KinesisProducer(config);
    return producer;
}
 

@SuppressWarnings("unchecked")
@Test
public void testRecordTooLarge() throws Exception {
  KinesisProducerConfigBean config = getKinesisTargetConfig();

  KinesisTarget target = new KinesisTarget(config, new ToOriginResponseConfig());
  TargetRunner targetRunner = new TargetRunner.Builder(
      KinesisDTarget.class,
      target
  ).setOnRecordError(OnRecordError.TO_ERROR).build();

  KinesisTestUtil.mockKinesisUtil(1);

  KinesisProducer producer = mock(KinesisProducer.class);
  Whitebox.setInternalState(target, "kinesisProducer", producer);

  targetRunner.runInit();

  ListenableFuture<UserRecordResult> future = mock(ListenableFuture.class);

  UserRecordResult result = mock(UserRecordResult.class);

  when(result.isSuccessful()).thenReturn(true);

  when(future.get()).thenReturn(result);

  when(producer.addUserRecord(any(String.class), any(String.class), any(ByteBuffer.class)))
      .thenReturn(future);

  List<Record> records = new ArrayList<>(4);
  records.add(KinesisTestUtil.getTooLargeRecord());
  records.addAll(KinesisTestUtil.getProducerTestRecords(3));
  targetRunner.runWrite(records);

  // Verify we added 3 good records at the end of the batch but not the bad one
  verify(producer, times(3)).addUserRecord(eq(STREAM_NAME), any(String.class), any(ByteBuffer.class));

  assertEquals(1, targetRunner.getErrorRecords().size());
  targetRunner.runDestroy();
}
 

/**
 * Given an AWS region, generate a configured Kinesis Producer.
 */
private static KinesisProducer getKinesisProducer(String region)
{
    KinesisProducerConfiguration config = new KinesisProducerConfiguration();
    config.setRegion(region);
    config.setMaxConnections(1);
    config.setRequestTimeout(60000);
    config.setRecordMaxBufferedTime(15000);
    config.setAggregationEnabled(true);

    return new KinesisProducer(config);
}
 

public KinesisEventProducer(String format, CommandLine cmd,
                            String propertiesFile, String streamName,
                            int throttleQueueSize, String appName) {
    this(throttleQueueSize);
    KinesisProducerConfiguration config = KinesisProducerConfiguration.fromPropertiesFile(propertiesFile);
    this.streamName = streamName;
    this.producer = new KinesisProducer(config);
    this.metrics = new KouplerMetrics(this, config, appName);
    this.format = FormatFactory.getFormat(format, cmd);
}
 

@Override
public void open(Map<String, Object> config, SinkContext sinkContext) throws Exception {
    kinesisSinkConfig = KinesisSinkConfig.load(config);
    this.sinkContext = sinkContext;

    checkArgument(isNotBlank(kinesisSinkConfig.getAwsKinesisStreamName()), "empty kinesis-stream name");
    checkArgument(isNotBlank(kinesisSinkConfig.getAwsEndpoint()) || 
                  isNotBlank(kinesisSinkConfig.getAwsRegion()), 
                  "Either the aws-end-point or aws-region must be set");
    checkArgument(isNotBlank(kinesisSinkConfig.getAwsCredentialPluginParam()), "empty aws-credential param");

    KinesisProducerConfiguration kinesisConfig = new KinesisProducerConfiguration();
    kinesisConfig.setKinesisEndpoint(kinesisSinkConfig.getAwsEndpoint());
    kinesisConfig.setRegion(kinesisSinkConfig.getAwsRegion());
    kinesisConfig.setThreadingModel(ThreadingModel.POOLED);
    kinesisConfig.setThreadPoolSize(4);
    kinesisConfig.setCollectionMaxCount(1);
    AWSCredentialsProvider credentialsProvider = createCredentialProvider(
            kinesisSinkConfig.getAwsCredentialPluginName(),
            kinesisSinkConfig.getAwsCredentialPluginParam())
        .getCredentialProvider();
    kinesisConfig.setCredentialsProvider(credentialsProvider);

    this.streamName = kinesisSinkConfig.getAwsKinesisStreamName();
    this.kinesisProducer = new KinesisProducer(kinesisConfig);
    IS_PUBLISH_FAILED.set(this, FALSE);

    LOG.info("Kinesis sink started. {}", (ReflectionToStringBuilder.toString(kinesisConfig, ToStringStyle.SHORT_PREFIX_STYLE)));
}
 

@Override
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return mockProducer;
}
 

/**
 * Creates a {@link KinesisProducer}.
 * Exposed so that tests can inject mock producers easily.
 */
@VisibleForTesting
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return new KinesisProducer(producerConfig);
}
 

/**
 * Creates a {@link KinesisProducer}.
 * Exposed so that tests can inject mock producers easily.
 */
@VisibleForTesting
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return new KinesisProducer(producerConfig);
}
 

public KPLClickEventsToKinesis(BlockingQueue<ClickEvent> inputQueue) {
    super(inputQueue);
    kinesis = new KinesisProducer(new KinesisProducerConfiguration()
            .setRegion(REGION)
            .setRecordMaxBufferedTime(5000));
}
 

protected AdvancedKPLClickEventsToKinesis(
        BlockingQueue<ClickEvent> inputQueue) {
    super(inputQueue);
    kinesis = new KinesisProducer(new KinesisProducerConfiguration()
            .setRegion(REGION));
}
 

@SuppressWarnings("unchecked")
@Test
public void testInOrderProduce() throws Exception {
  KinesisProducerConfigBean config = getKinesisTargetConfig();
  config.preserveOrdering = true;

  KinesisTarget target = new KinesisTarget(config, new ToOriginResponseConfig());
  TargetRunner targetRunner = new TargetRunner.Builder(KinesisDTarget.class, target).build();

  PowerMockito.mockStatic(KinesisUtil.class);

  when(KinesisUtil.checkStreamExists(
      any(ClientConfiguration.class),
      any(KinesisConfigBean.class),
      any(String.class),
      any(List.class),
      any(Stage.Context.class)
      )
  ).thenReturn(1L);

  KinesisProducer producer = mock(KinesisProducer.class);
  Whitebox.setInternalState(target, "kinesisProducer", producer);

  targetRunner.runInit();

  ListenableFuture<UserRecordResult> future = mock(ListenableFuture.class);

  UserRecordResult result = mock(UserRecordResult.class);

  when(result.isSuccessful()).thenReturn(true);
  when(result.getShardId()).thenReturn("shardId-000000000000");

  when(future.get()).thenReturn(result);

  when(producer.addUserRecord(any(String.class), any(String.class), any(ByteBuffer.class)))
      .thenReturn(future);

  targetRunner.runWrite(KinesisTestUtil.getProducerTestRecords(3));

  // Verify we added 3 records to stream test
  verify(producer, times(3)).addUserRecord(eq(STREAM_NAME), any(String.class), any(ByteBuffer.class));
  // With preserveOrdering we should call flushSync for each record, plus once more for the batch.
  // The last invocation has no effect as no records should be pending.
  verify(producer, times(4)).flushSync();

  targetRunner.runDestroy();
}
 

@Override
public IKinesisProducer createKinesisProducer(KinesisProducerConfiguration config) {
  config.setRegion(region.getName());
  config.setCredentialsProvider(getCredentialsProvider());
  return new KinesisProducer(config);
}
 

public static void main(String[] args) throws Exception
{
    if (args.length != 2)
    {
        System.err.println("Usage SampleKPLProducer <stream name> <region>");
        System.exit(1);
    }

    String streamName = args[0];
    String regionName = args[1];

    final KinesisProducer producer = getKinesisProducer(regionName);

    final AtomicLong sequenceNumber = new AtomicLong(0);
    final AtomicLong completed = new AtomicLong(0);

    final FutureCallback<UserRecordResult> callback = new FutureCallback<UserRecordResult>()
    {
        @Override
        public void onFailure(Throwable t)
        {
            if (t instanceof UserRecordFailedException)
            {
                Attempt last = Iterables.getLast(((UserRecordFailedException) t).getResult().getAttempts());
                System.err.println(String.format("Record failed to put - %s : %s", last.getErrorCode(), last.getErrorMessage()));
            }
            System.err.println("Exception during put: " + t.getMessage());
            t.printStackTrace();
            System.exit(1);
        }

        @Override
        public void onSuccess(UserRecordResult result)
        {
            completed.getAndIncrement();
        }
    };

    final Runnable putOneRecord = new Runnable()
    {
        @Override
        public void run()
        {
            byte[] data = ProducerUtils.randomData(sequenceNumber.get(), ProducerConfig.RECORD_SIZE_BYTES);
            ListenableFuture<UserRecordResult> f = producer.addUserRecord(streamName, ProducerUtils.randomPartitionKey(),
                    ProducerUtils.randomExplicitHashKey(), ByteBuffer.wrap(data));
            Futures.addCallback(f, callback);
        }
    };

    EXECUTOR.scheduleAtFixedRate(new Runnable()
    {
        @Override
        public void run()
        {
            long put = sequenceNumber.get();
            long total = RECORDS_PER_SECOND * SECONDS_TO_RUN;
            double putPercent = 100.0 * put / total;
            long done = completed.get();
            double donePercent = 100.0 * done / total;
            System.out.println(String.format("Put %d of %d so far (%.2f %%), %d have completed (%.2f %%)", put, total, putPercent, done, donePercent));
        }
    }, 1, 1, TimeUnit.SECONDS);

    System.out.println(String.format("Starting puts... will run for %d seconds at %d records per second", SECONDS_TO_RUN, RECORDS_PER_SECOND));

    executeAtTargetRate(EXECUTOR, putOneRecord, sequenceNumber, SECONDS_TO_RUN, RECORDS_PER_SECOND);

    EXECUTOR.awaitTermination(SECONDS_TO_RUN + 1, TimeUnit.SECONDS);

    System.out.println("Waiting for remaining puts to finish...");
    producer.flushSync();
    System.out.println("All records complete.");

    producer.destroy();
    System.out.println("Finished.");
}
 

/** The main method.
 *  @param args  The command line args for the Sample Producer. It takes 3 optional position parameters:
 *  1. The stream name to use (default-data-stream is default)
 *  2. The region name to use (us-east-1 is default)
 *  3. The duration of the test in seconds, 5 is the default.
 */
public static void main(String[] args) throws Exception {
    final String streamName = getArgIfPresent(args, 0, STREAM_NAME);
    final String region = getArgIfPresent(args, 1, REGION);
    final String secondsToRunString = getArgIfPresent(args, 2, String.valueOf(SECONDS_TO_RUN_DEFAULT));
    final int secondsToRun = Integer.parseInt(secondsToRunString);
    if (secondsToRun <= 0) {
        LOGGER.error("Seconds to Run should be a positive integer");
        System.exit(1);
    }

    final KinesisProducer producer = getKinesisProducer(region);
    final AtomicLong sequenceNumber = new AtomicLong(0);
    final AtomicLong completed = new AtomicLong(0);

    FutureCallback<UserRecordResult> callback = new FutureCallback<UserRecordResult>() {
        @Override public void onFailure(Throwable t) {
            // If we see any failures, we will log them.
            if (t instanceof UserRecordFailedException) {
                Attempt last = Iterables.getLast(((UserRecordFailedException) t).getResult().getAttempts());
                LOGGER.error(String.format("Record failed to put - %s : %s", last.getErrorCode(), last.getErrorMessage()));
            }
            LOGGER.error("Exception during put", t);
        };

        @Override public void onSuccess(UserRecordResult result) {
            completed.getAndIncrement();
        };
    };

    final ExecutorService callbackThreadPool = Executors.newCachedThreadPool();

    // The lines within run() are the essence of the KPL API.
    final Runnable putOneRecord = new Runnable() {
        @Override
        public void run() {
            ByteBuffer data = generateData();
            // TIMESTAMP is our partition key
            ListenableFuture<UserRecordResult> f = producer.addUserRecord(streamName, TIMESTAMP, randomExplicitHashKey(), data);
            Futures.addCallback(f, callback, callbackThreadPool);
        }
    };

    EXECUTOR.scheduleAtFixedRate(new Runnable() {
        @Override
        public void run() {
            long put = sequenceNumber.get();
            long total = RECORDS_PER_SECOND * secondsToRun;
            double putPercent = 100.0 * put / total;
            long done = completed.get();
            double donePercent = 100.0 * done / total;
            LOGGER.info(String.format(
                "Put %d of %d so far (%.2f %%), %d have completed (%.2f %%)",
                put, total, putPercent, done, donePercent
            ));
        }
    }, 1, 1, TimeUnit.SECONDS);

    executeAtTargetRate(EXECUTOR, putOneRecord, sequenceNumber, secondsToRun, RECORDS_PER_SECOND);

    EXECUTOR.awaitTermination(secondsToRun + 1, TimeUnit.SECONDS);

    LOGGER.info("Waiting for remaining puts to finish...");
    producer.flushSync();
    LOGGER.info("All records complete.");

    producer.destroy();
    LOGGER.info("Finished.");
}
 

public Stream(KinesisProducer kinesisProducer, String streamName) {
    this.kinesisProducer = kinesisProducer;
    this.streamName = streamName;
}
 

@Override
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return mockProducer;
}
 

/**
 * Creates a {@link KinesisProducer}.
 * Exposed so that tests can inject mock producers easily.
 */
@VisibleForTesting
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return new KinesisProducer(producerConfig);
}
 

@Override
protected KinesisProducer getKinesisProducer(KinesisProducerConfiguration producerConfig) {
	return mockProducer;
}
 

@SuppressWarnings("unchecked")
@Test
public void testDefaultProduce() throws Exception {
  KinesisProducerConfigBean config = getKinesisTargetConfig();

  KinesisTarget target = new KinesisTarget(config, new ToOriginResponseConfig());
  TargetRunner targetRunner = new TargetRunner.Builder(KinesisDTarget.class, target).build();

  KinesisTestUtil.mockKinesisUtil(1);

  KinesisProducer producer = mock(KinesisProducer.class);
  Whitebox.setInternalState(target, "kinesisProducer", producer);

  targetRunner.runInit();

  ListenableFuture<UserRecordResult> future = mock(ListenableFuture.class);

  UserRecordResult result = mock(UserRecordResult.class);

  when(result.isSuccessful()).thenReturn(true);

  when(future.get()).thenReturn(result);

  when(producer.addUserRecord(any(String.class), any(String.class), any(ByteBuffer.class)))
      .thenReturn(future);

  targetRunner.runWrite(KinesisTestUtil.getProducerTestRecords(3));

  // Verify we added 3 records
  verify(producer, times(3)).addUserRecord(eq(STREAM_NAME), any(String.class), any(ByteBuffer.class));
  // By default we should only call flushSync one time per batch.
  verify(producer, times(1)).flushSync();

  targetRunner.runDestroy();
}