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

private Record generateDestResponseRecord(UserRecordResult result) {
  Record record = getContext().createRecord("responseRecord");
  List<Field> attemptsVal = new ArrayList<>();
  for(Attempt attempt: result.getAttempts()) {
    LinkedHashMap<String, Field> attemptVal = new LinkedHashMap<>();
    attemptVal.put("delay", Field.create(attempt.getDelay()));
    attemptVal.put("duration", Field.create(attempt.getDuration()));
    attemptVal.put("errorMessage", Field.create(attempt.getErrorMessage()));
    attemptVal.put("errorCode", Field.create(attempt.getErrorCode()));
    attemptVal.put("success", Field.create(attempt.isSuccessful()));
    attemptsVal.add(Field.createListMap(attemptVal));
  }
  LinkedHashMap<String, Field> resultVal = new LinkedHashMap<>();
  resultVal.put("sequenceNumber", Field.create(result.getSequenceNumber()));
  resultVal.put("shardId", Field.create(result.getShardId()));
  resultVal.put("successful", Field.create(result.isSuccessful()));
  resultVal.put("attempts", Field.create(attemptsVal));
  record.set(Field.createListMap(resultVal));
  return record;
}
 

/**
 * It adds a record asynchronously which then should be delivered by Kinesis producer in
 * background (Kinesis producer forks native processes to do this job).
 *
 * <p>The records can be batched and then they will be sent in one HTTP request. Amazon KPL
 * supports two types of batching - aggregation and collection - and they can be configured by
 * producer properties.
 *
 * <p>More details can be found here: <a
 * href="https://docs.aws.amazon.com/streams/latest/dev/kinesis-kpl-concepts.html">KPL Key
 * Concepts</a> and <a
 * href="https://docs.aws.amazon.com/streams/latest/dev/kinesis-kpl-config.html">Configuring
 * the KPL</a>
 */
@ProcessElement
public void processElement(ProcessContext c) {
  ByteBuffer data = ByteBuffer.wrap(c.element());
  String partitionKey = spec.getPartitionKey();
  String explicitHashKey = null;

  // Use custom partitioner
  if (partitioner != null) {
    partitionKey = partitioner.getPartitionKey(c.element());
    explicitHashKey = partitioner.getExplicitHashKey(c.element());
  }

  ListenableFuture<UserRecordResult> f =
      producer.addUserRecord(spec.getStreamName(), partitionKey, explicitHashKey, data);
  putFutures.add(f);
}
 

@Override
public void write(Record<byte[]> record) throws Exception {
    // kpl-thread captures publish-failure. fail the publish on main pulsar-io-thread to maintain the ordering 
    if (kinesisSinkConfig.isRetainOrdering() && previousPublishFailed == TRUE) {
        LOG.warn("Skip acking message to retain ordering with previous failed message {}-{}", this.streamName,
                record.getRecordSequence());
        throw new IllegalStateException("kinesis queue has publish failure");
    }
    String partitionedKey = record.getKey().orElse(record.getTopicName().orElse(defaultPartitionedKey));
    partitionedKey = partitionedKey.length() > maxPartitionedKeyLength
            ? partitionedKey.substring(0, maxPartitionedKeyLength - 1)
            : partitionedKey; // partitionedKey Length must be at least one, and at most 256
    ByteBuffer data = createKinesisMessage(kinesisSinkConfig.getMessageFormat(), record);
    ListenableFuture<UserRecordResult> addRecordResult = kinesisProducer.addUserRecord(this.streamName,
            partitionedKey, data);
    addCallback(addRecordResult,
            ProducerSendCallback.create(this, record, System.nanoTime()), directExecutor());
    if (sinkContext != null) {
        sinkContext.recordMetric(METRICS_TOTAL_INCOMING, 1);
        sinkContext.recordMetric(METRICS_TOTAL_INCOMING_BYTES, data.array().length);
    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("Published message to kinesis stream {} with size {}", streamName, record.getValue().length);
    }
}
 

@Override
public void onSuccess(UserRecordResult result) {
    if (LOG.isDebugEnabled()) {
        LOG.debug("Successfully published message for {}-{} with latency {}", kinesisSink.streamName, result.getShardId(),
                TimeUnit.NANOSECONDS.toMillis((System.nanoTime() - startTime)));
    }
    if (kinesisSink.sinkContext != null) {
        kinesisSink.sinkContext.recordMetric(METRICS_TOTAL_SUCCESS, 1);
    }
    if (kinesisSink.kinesisSinkConfig.isRetainOrdering() && kinesisSink.previousPublishFailed == TRUE) {
        LOG.warn("Skip acking message to retain ordering with previous failed message {}-{} on shard {}",
                kinesisSink.streamName, resultContext.getRecordSequence(), result.getShardId());
    } else {
        this.resultContext.ack();
    }
    recycle();
}
 

public void send(String event) throws UnsupportedEncodingException {
    byte[] bytes = event.getBytes("UTF-8");
    this.metrics.queueEvent(bytes.length);
    ByteBuffer data = ByteBuffer.wrap(bytes);
    String partitionKey = getPartitionKey(event);
    if (partitionKey != null) {
        ListenableFuture<UserRecordResult> f = producer.addUserRecord(streamName, partitionKey, data);
        Futures.addCallback(f, new FutureCallback<UserRecordResult>() {
            @Override
            public void onFailure(Throwable t) {
                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) {
                metrics.ackEvent();
            }
        });
    }
}
 

private void getAndCheck(Future<UserRecordResult> future) throws StageException {
  try {
    UserRecordResult result = future.get();
    if (!result.isSuccessful()) {
      for (Attempt attempt : result.getAttempts()) {
        LOG.error("Failed to put record: {}", attempt.getErrorMessage());
      }
      throw new StageException(Errors.KINESIS_00, result.getAttempts().get(0).getErrorMessage());
    }

    if (responseConf.sendResponseToOrigin &&
        this.responseConf.responseType.equals(ResponseType.DESTINATION_RESPONSE)) {
      getContext().toSourceResponse(generateDestResponseRecord(result));
    }

  } catch (InterruptedException | ExecutionException e) {
    LOG.error("Pipeline is shutting down.", e);
    // We should flush if we encounter an error.
    kinesisProducer.flushSync();
  }
}
 

private int getNumPendingRecordFutures() {
	int numPending = 0;

	for (SettableFuture<UserRecordResult> future : pendingRecordFutures) {
		if (!future.isDone()) {
			numPending++;
		}
	}

	return numPending;
}
 

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

@Override
public synchronized ListenableFuture<UserRecordResult> addUserRecord(
    String stream, String partitionKey, String explicitHashKey, ByteBuffer data) {
  seqNumber.incrementAndGet();
  SettableFuture<UserRecordResult> f = SettableFuture.create();
  f.set(
      new UserRecordResult(
          new ArrayList<>(), String.valueOf(seqNumber.get()), explicitHashKey, !isFailedFlush));

  if (kinesisService.getExistedStream().equals(stream)) {
    addedRecords.add(new UserRecord(stream, partitionKey, explicitHashKey, data));
  }
  return f;
}
 

private int getNumPendingRecordFutures() {
	int numPending = 0;

	for (SettableFuture<UserRecordResult> future : pendingRecordFutures) {
		if (!future.isDone()) {
			numPending++;
		}
	}

	return numPending;
}
 

private int getNumPendingRecordFutures() {
	int numPending = 0;

	for (SettableFuture<UserRecordResult> future : pendingRecordFutures) {
		if (!future.isDone()) {
			numPending++;
		}
	}

	return numPending;
}
 

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 put(Collection<SinkRecord> sinkRecords) {

	// If KinesisProducers cannot write to Kinesis Streams (because of
	// connectivity issues, access issues
	// or misconfigured shards we will pause consumption of messages till
	// backlog is cleared

	validateOutStandingRecords();

	String partitionKey;
	for (SinkRecord sinkRecord : sinkRecords) {

		ListenableFuture<UserRecordResult> f;
		// Kinesis does not allow empty partition key
		if (sinkRecord.key() != null && !sinkRecord.key().toString().trim().equals("")) {
			partitionKey = sinkRecord.key().toString().trim();
		} else {
			partitionKey = Integer.toString(sinkRecord.kafkaPartition());
		}

		if (singleKinesisProducerPerPartition)
			f = addUserRecord(producerMap.get(sinkRecord.kafkaPartition() + "@" + sinkRecord.topic()), streamName,
					partitionKey, usePartitionAsHashKey, sinkRecord);
		else
			f = addUserRecord(kinesisProducer, streamName, partitionKey, usePartitionAsHashKey, sinkRecord);

		Futures.addCallback(f, callback);

	}
}
 

/**
 * Test ensuring that the producer is not dropping buffered records;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@SuppressWarnings({"unchecked", "ResultOfMethodCallIgnored"})
@Test(timeout = 10000)
public void testAtLeastOnceProducer() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));
	testHarness.processElement(new StreamRecord<>("msg-2"));
	testHarness.processElement(new StreamRecord<>("msg-3"));

	// start a thread to perform checkpointing
	CheckedThread snapshotThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block until all records are flushed;
			// if the snapshot implementation returns before pending records are flushed,
			testHarness.snapshot(123L, 123L);
		}
	};
	snapshotThread.start();

	// before proceeding, make sure that flushing has started and that the snapshot is still blocked;
	// this would block forever if the snapshot didn't perform a flush
	producer.waitUntilFlushStarted();
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	// now, complete the callbacks
	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);

	producer.getPendingRecordFutures().get(0).set(result);
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	producer.getPendingRecordFutures().get(1).set(result);
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	producer.getPendingRecordFutures().get(2).set(result);

	// this would fail with an exception if flushing wasn't completed before the snapshot method returned
	snapshotThread.sync();

	testHarness.close();
}
 

@Override
public ListenableFuture<UserRecordResult> addUserRecord(
    String stream, String partitionKey, ByteBuffer data) {
  throw new UnsupportedOperationException("Not implemented");
}
 

@Override
public ListenableFuture<UserRecordResult> addUserRecord(UserRecord userRecord) {
  throw new UnsupportedOperationException("Not implemented");
}
 

List<SettableFuture<UserRecordResult>> getPendingRecordFutures() {
	return pendingRecordFutures;
}
 

/**
 * Test ensuring that the producer blocks if the queue limit is exceeded,
 * until the queue length drops below the limit;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@Test(timeout = 10000)
public void testBackpressure() throws Throwable {
	final Deadline deadline = Deadline.fromNow(Duration.ofSeconds(10));

	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());
	producer.setQueueLimit(1);

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
			new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);

	CheckedThread msg1 = new CheckedThread() {
		@Override
		public void go() throws Exception {
			testHarness.processElement(new StreamRecord<>("msg-1"));
		}
	};
	msg1.start();
	msg1.trySync(deadline.timeLeftIfAny().toMillis());
	assertFalse("Flush triggered before reaching queue limit", msg1.isAlive());

	// consume msg-1 so that queue is empty again
	producer.getPendingRecordFutures().get(0).set(result);

	CheckedThread msg2 = new CheckedThread() {
		@Override
		public void go() throws Exception {
			testHarness.processElement(new StreamRecord<>("msg-2"));
		}
	};
	msg2.start();
	msg2.trySync(deadline.timeLeftIfAny().toMillis());
	assertFalse("Flush triggered before reaching queue limit", msg2.isAlive());

	CheckedThread moreElementsThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block until msg-2 is consumed
			testHarness.processElement(new StreamRecord<>("msg-3"));
			// this should block until msg-3 is consumed
			testHarness.processElement(new StreamRecord<>("msg-4"));
		}
	};
	moreElementsThread.start();

	assertTrue("Producer should still block, but doesn't", moreElementsThread.isAlive());

	// consume msg-2 from the queue, leaving msg-3 in the queue and msg-4 blocked
	while (producer.getPendingRecordFutures().size() < 2) {
		Thread.sleep(50);
	}
	producer.getPendingRecordFutures().get(1).set(result);

	assertTrue("Producer should still block, but doesn't", moreElementsThread.isAlive());

	// consume msg-3, blocked msg-4 can be inserted into the queue and block is released
	while (producer.getPendingRecordFutures().size() < 3) {
		Thread.sleep(50);
	}
	producer.getPendingRecordFutures().get(2).set(result);

	moreElementsThread.trySync(deadline.timeLeftIfAny().toMillis());

	assertFalse("Prodcuer still blocks although the queue is flushed", moreElementsThread.isAlive());

	producer.getPendingRecordFutures().get(3).set(result);

	testHarness.close();
}
 

/**
 * Test ensuring that if an async exception is caught for one of the flushed requests on checkpoint,
 * it should be rethrown; we set a timeout because the test will not finish if the logic is broken.
 *
 * <p>Note that this test does not test the snapshot method is blocked correctly when there are pending records.
 * The test for that is covered in testAtLeastOnceProducer.
 */
@SuppressWarnings("ResultOfMethodCallIgnored")
@Test(timeout = 10000)
public void testAsyncErrorRethrownAfterFlush() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));
	testHarness.processElement(new StreamRecord<>("msg-2"));
	testHarness.processElement(new StreamRecord<>("msg-3"));

	// only let the first record succeed for now
	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);
	producer.getPendingRecordFutures().get(0).set(result);

	CheckedThread snapshotThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block at first, since there are still two pending records that needs to be flushed
			testHarness.snapshot(123L, 123L);
		}
	};
	snapshotThread.start();

	// let the 2nd message fail with an async exception
	producer.getPendingRecordFutures().get(1).setException(new Exception("artificial async failure for 2nd message"));
	producer.getPendingRecordFutures().get(2).set(mock(UserRecordResult.class));

	try {
		snapshotThread.sync();
	} catch (Exception e) {
		// after the flush, the async exception should have been rethrown
		Assert.assertTrue(ExceptionUtils.findThrowableWithMessage(e, "artificial async failure for 2nd message").isPresent());

		// test succeeded
		return;
	}

	Assert.fail();
}
 

/** 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.");
}
 

@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 void open(Configuration parameters) throws Exception {
	super.open(parameters);

	schema.open(() -> getRuntimeContext().getMetricGroup().addGroup("user"));

	// check and pass the configuration properties
	KinesisProducerConfiguration producerConfig = KinesisConfigUtil.getValidatedProducerConfiguration(configProps);

	producer = getKinesisProducer(producerConfig);

	final MetricGroup kinesisMectricGroup = getRuntimeContext().getMetricGroup().addGroup(KINESIS_PRODUCER_METRIC_GROUP);
	this.backpressureCycles = kinesisMectricGroup.counter(METRIC_BACKPRESSURE_CYCLES);
	kinesisMectricGroup.gauge(METRIC_OUTSTANDING_RECORDS_COUNT, producer::getOutstandingRecordsCount);

	backpressureLatch = new TimeoutLatch();
	callback = new FutureCallback<UserRecordResult>() {
		@Override
		public void onSuccess(UserRecordResult result) {
			backpressureLatch.trigger();
			if (!result.isSuccessful()) {
				if (failOnError) {
					// only remember the first thrown exception
					if (thrownException == null) {
						thrownException = new RuntimeException("Record was not sent successful");
					}
				} else {
					LOG.warn("Record was not sent successful");
				}
			}
		}

		@Override
		public void onFailure(Throwable t) {
			backpressureLatch.trigger();
			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 invoke(OUT value, Context context) throws Exception {
	if (this.producer == null) {
		throw new RuntimeException("Kinesis producer has been closed");
	}

	checkAndPropagateAsyncError();
	boolean didWaitForFlush = enforceQueueLimit();

	if (didWaitForFlush) {
		checkAndPropagateAsyncError();
	}

	String stream = defaultStream;
	String partition = defaultPartition;

	ByteBuffer serialized = schema.serialize(value);

	// maybe set custom stream
	String customStream = schema.getTargetStream(value);
	if (customStream != null) {
		stream = customStream;
	}

	String explicitHashkey = null;
	// maybe set custom partition
	if (customPartitioner != null) {
		partition = customPartitioner.getPartitionId(value);
		explicitHashkey = customPartitioner.getExplicitHashKey(value);
	}

	if (stream == null) {
		if (failOnError) {
			throw new RuntimeException("No target stream set");
		} else {
			LOG.warn("No target stream set. Skipping record");
			return;
		}
	}

	ListenableFuture<UserRecordResult> cb = producer.addUserRecord(stream, partition, explicitHashkey, serialized);
	Futures.addCallback(cb, callback);
}
 

@Override
public void open(Configuration parameters) throws Exception {
	super.open(parameters);

	// check and pass the configuration properties
	KinesisProducerConfiguration producerConfig = KinesisConfigUtil.getValidatedProducerConfiguration(configProps);

	producer = getKinesisProducer(producerConfig);

	final MetricGroup kinesisMectricGroup = getRuntimeContext().getMetricGroup().addGroup(KINESIS_PRODUCER_METRIC_GROUP);
	this.backpressureCycles = kinesisMectricGroup.counter(METRIC_BACKPRESSURE_CYCLES);
	kinesisMectricGroup.gauge(METRIC_OUTSTANDING_RECORDS_COUNT, producer::getOutstandingRecordsCount);

	backpressureLatch = new TimeoutLatch();
	callback = new FutureCallback<UserRecordResult>() {
		@Override
		public void onSuccess(UserRecordResult result) {
			backpressureLatch.trigger();
			if (!result.isSuccessful()) {
				if (failOnError) {
					// only remember the first thrown exception
					if (thrownException == null) {
						thrownException = new RuntimeException("Record was not sent successful");
					}
				} else {
					LOG.warn("Record was not sent successful");
				}
			}
		}

		@Override
		public void onFailure(Throwable t) {
			backpressureLatch.trigger();
			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 invoke(OUT value, Context context) throws Exception {
	if (this.producer == null) {
		throw new RuntimeException("Kinesis producer has been closed");
	}

	checkAndPropagateAsyncError();
	boolean didWaitForFlush = enforceQueueLimit();

	if (didWaitForFlush) {
		checkAndPropagateAsyncError();
	}

	String stream = defaultStream;
	String partition = defaultPartition;

	ByteBuffer serialized = schema.serialize(value);

	// maybe set custom stream
	String customStream = schema.getTargetStream(value);
	if (customStream != null) {
		stream = customStream;
	}

	String explicitHashkey = null;
	// maybe set custom partition
	if (customPartitioner != null) {
		partition = customPartitioner.getPartitionId(value);
		explicitHashkey = customPartitioner.getExplicitHashKey(value);
	}

	if (stream == null) {
		if (failOnError) {
			throw new RuntimeException("No target stream set");
		} else {
			LOG.warn("No target stream set. Skipping record");
			return;
		}
	}

	ListenableFuture<UserRecordResult> cb = producer.addUserRecord(stream, partition, explicitHashkey, serialized);
	Futures.addCallback(cb, callback);
}
 

/**
 * Test ensuring that if an async exception is caught for one of the flushed requests on checkpoint,
 * it should be rethrown; we set a timeout because the test will not finish if the logic is broken.
 *
 * <p>Note that this test does not test the snapshot method is blocked correctly when there are pending records.
 * The test for that is covered in testAtLeastOnceProducer.
 */
@SuppressWarnings("ResultOfMethodCallIgnored")
@Test(timeout = 10000)
public void testAsyncErrorRethrownAfterFlush() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));
	testHarness.processElement(new StreamRecord<>("msg-2"));
	testHarness.processElement(new StreamRecord<>("msg-3"));

	// only let the first record succeed for now
	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);
	producer.getPendingRecordFutures().get(0).set(result);

	CheckedThread snapshotThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block at first, since there are still two pending records that needs to be flushed
			testHarness.snapshot(123L, 123L);
		}
	};
	snapshotThread.start();

	// let the 2nd message fail with an async exception
	producer.getPendingRecordFutures().get(1).setException(new Exception("artificial async failure for 2nd message"));
	producer.getPendingRecordFutures().get(2).set(mock(UserRecordResult.class));

	try {
		snapshotThread.sync();
	} catch (Exception e) {
		// after the flush, the async exception should have been rethrown
		Assert.assertTrue(ExceptionUtils.findThrowableWithMessage(e, "artificial async failure for 2nd message").isPresent());

		// test succeeded
		return;
	}

	Assert.fail();
}
 

/**
 * Test ensuring that the producer is not dropping buffered records;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@SuppressWarnings({"unchecked", "ResultOfMethodCallIgnored"})
@Test(timeout = 10000)
public void testAtLeastOnceProducer() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));
	testHarness.processElement(new StreamRecord<>("msg-2"));
	testHarness.processElement(new StreamRecord<>("msg-3"));

	// start a thread to perform checkpointing
	CheckedThread snapshotThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block until all records are flushed;
			// if the snapshot implementation returns before pending records are flushed,
			testHarness.snapshot(123L, 123L);
		}
	};
	snapshotThread.start();

	// before proceeding, make sure that flushing has started and that the snapshot is still blocked;
	// this would block forever if the snapshot didn't perform a flush
	producer.waitUntilFlushStarted();
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	// now, complete the callbacks
	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);

	producer.getPendingRecordFutures().get(0).set(result);
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	producer.getPendingRecordFutures().get(1).set(result);
	Assert.assertTrue("Snapshot returned before all records were flushed", snapshotThread.isAlive());

	producer.getPendingRecordFutures().get(2).set(result);

	// this would fail with an exception if flushing wasn't completed before the snapshot method returned
	snapshotThread.sync();

	testHarness.close();
}
 

/**
 * Test ensuring that the producer blocks if the queue limit is exceeded,
 * until the queue length drops below the limit;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@Test(timeout = 10000)
public void testBackpressure() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());
	producer.setQueueLimit(1);

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
			new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	UserRecordResult result = mock(UserRecordResult.class);
	when(result.isSuccessful()).thenReturn(true);

	CheckedThread msg1 = new CheckedThread() {
		@Override
		public void go() throws Exception {
			testHarness.processElement(new StreamRecord<>("msg-1"));
		}
	};
	msg1.start();
	msg1.trySync(100);
	assertFalse("Flush triggered before reaching queue limit", msg1.isAlive());

	// consume msg-1 so that queue is empty again
	producer.getPendingRecordFutures().get(0).set(result);

	CheckedThread msg2 = new CheckedThread() {
		@Override
		public void go() throws Exception {
			testHarness.processElement(new StreamRecord<>("msg-2"));
		}
	};
	msg2.start();
	msg2.trySync(100);
	assertFalse("Flush triggered before reaching queue limit", msg2.isAlive());

	CheckedThread moreElementsThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			// this should block until msg-2 is consumed
			testHarness.processElement(new StreamRecord<>("msg-3"));
			// this should block until msg-3 is consumed
			testHarness.processElement(new StreamRecord<>("msg-4"));
		}
	};
	moreElementsThread.start();

	moreElementsThread.trySync(100);
	assertTrue("Producer should still block, but doesn't", moreElementsThread.isAlive());

	// consume msg-2 from the queue, leaving msg-3 in the queue and msg-4 blocked
	producer.getPendingRecordFutures().get(1).set(result);

	moreElementsThread.trySync(100);
	assertTrue("Producer should still block, but doesn't", moreElementsThread.isAlive());

	// consume msg-3, blocked msg-4 can be inserted into the queue and block is released
	producer.getPendingRecordFutures().get(2).set(result);

	moreElementsThread.trySync(100);

	assertFalse("Prodcuer still blocks although the queue is flushed", moreElementsThread.isAlive());

	producer.getPendingRecordFutures().get(3).set(result);

	testHarness.close();
}
 

List<SettableFuture<UserRecordResult>> getPendingRecordFutures() {
	return pendingRecordFutures;
}
 

@Override
public void open(Configuration parameters) throws Exception {
	super.open(parameters);

	// check and pass the configuration properties
	KinesisProducerConfiguration producerConfig = KinesisConfigUtil.getValidatedProducerConfiguration(configProps);

	producer = getKinesisProducer(producerConfig);

	final MetricGroup kinesisMectricGroup = getRuntimeContext().getMetricGroup().addGroup(KINESIS_PRODUCER_METRIC_GROUP);
	this.backpressureCycles = kinesisMectricGroup.counter(METRIC_BACKPRESSURE_CYCLES);
	kinesisMectricGroup.gauge(METRIC_OUTSTANDING_RECORDS_COUNT, producer::getOutstandingRecordsCount);

	backpressureLatch = new TimeoutLatch();
	callback = new FutureCallback<UserRecordResult>() {
		@Override
		public void onSuccess(UserRecordResult result) {
			backpressureLatch.trigger();
			if (!result.isSuccessful()) {
				if (failOnError) {
					// only remember the first thrown exception
					if (thrownException == null) {
						thrownException = new RuntimeException("Record was not sent successful");
					}
				} else {
					LOG.warn("Record was not sent successful");
				}
			}
		}

		@Override
		public void onFailure(Throwable t) {
			backpressureLatch.trigger();
			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());
}