SourceFunction
@Public
public interface SourceFunction<T> extends Function, Serializable {
void run(SourceContext<T> ctx) throws Exception;
void cancel();
interface SourceContext<T> {
void collect(T element);
@PublicEvolving
void collectWithTimestamp(T element, long timestamp);
void emitWatermark(Watermark mark);
@PublicEvolving
void markAsTemporarilyIdle();
Object getCheckpointLock();
void close();
}
}SourceFunction中还嵌套定义了SourceContext接口，它表示这个Source对应的上下文，用来发射数据。其中起主要作用的是前三个方法：
collect()：发射一个不带自定义时间戳的元素。如果流程序的时间特征（TimeCharacteristic）是处理时间（ProcessingTime），元素没有时间戳；如果是摄入时间（IngestionTime），元素会附带系统时间；如果是事件时间（EventTime），那么初始没有时间戳，但一旦要做与时间戳相关的操作（如窗口）时，就必须用TimestampAssigner设定一个。
collectWithTimestamp()：发射一个带有自定义时间戳的元素。该方法对于时间特征为事件时间的程序是绝对必须的，如果为处理时间就会被直接忽略，如果为摄入时间就会被系统时间覆盖。
emitWatermark()：发射一个水印，仅对于事件时间有效。一个带有时间戳t的水印表示不会有任何t' <= t的事件再发生，如果发生，会被当做迟到事件忽略掉。
ParallelSourceFunction，表示该Source可以按照设置的并行度并发执行。
RichSourceFunction，继承自富函数RichFunction，表示该Source可以感知到运行时上下文（RuntimeContext，如Task、State、并行度的信息），以及可以自定义初始化和销毁逻辑（通过open()/close()方法）。
RichParallelSourceFunction，以上两者的综合。
RocketMQ Source
public class RocketMQSource<OUT> extends RichParallelSourceFunction<OUT>
implements CheckpointedFunction, CheckpointListener, ResultTypeQueryable<OUT> {
private static final long serialVersionUID = 1L;
private static final Logger LOG = LoggerFactory.getLogger(RocketMQSource.class);
private transient MQPullConsumerScheduleService pullConsumerScheduleService;
private DefaultMQPullConsumer consumer;
private KeyValueDeserializationSchema<OUT> schema;
private RunningChecker runningChecker;
private transient ListState<Tuple2<MessageQueue, Long>> unionOffsetStates;
private Map<MessageQueue, Long> offsetTable;
private Map<MessageQueue, Long> restoredOffsets;
private LinkedMap pendingOffsetsToCommit;
private Properties props;
private String topic;
private String group;
private static final String OFFSETS_STATE_NAME = "topic-partition-offset-states";
private transient volatile boolean restored;
private transient boolean enableCheckpoint;
public RocketMQSource(KeyValueDeserializationSchema<OUT> schema, Properties props) {
this.schema = schema;
this.props = props;
}
@Override
public void open(Configuration parameters) throws Exception {
LOG.debug("source open....");
Validate.notEmpty(props, "Consumer properties can not be empty");
Validate.notNull(schema, "KeyValueDeserializationSchema can not be null");
this.topic = props.getProperty(RocketMQConfig.CONSUMER_TOPIC);
this.group = props.getProperty(RocketMQConfig.CONSUMER_GROUP);
Validate.notEmpty(topic, "Consumer topic can not be empty");
Validate.notEmpty(group, "Consumer group can not be empty");
this.enableCheckpoint = ((StreamingRuntimeContext) getRuntimeContext()).isCheckpointingEnabled();
if (offsetTable == null) {
offsetTable = new ConcurrentHashMap<>();
}
if (restoredOffsets == null) {
restoredOffsets = new ConcurrentHashMap<>();
}
if (pendingOffsetsToCommit == null) {
pendingOffsetsToCommit = new LinkedMap();
}
runningChecker = new RunningChecker();
pullConsumerScheduleService = new MQPullConsumerScheduleService(group);
consumer = pullConsumerScheduleService.getDefaultMQPullConsumer();
consumer.setInstanceName(String.valueOf(getRuntimeContext().getIndexOfThisSubtask()) + "_" + UUID.randomUUID());
RocketMQConfig.buildConsumerConfigs(props, consumer);
}
@Override
public void run(SourceContext context) throws Exception {
LOG.debug("source run....");
final Object lock = context.getCheckpointLock();
int delayWhenMessageNotFound = getInteger(props, RocketMQConfig.CONSUMER_DELAY_WHEN_MESSAGE_NOT_FOUND,
RocketMQConfig.DEFAULT_CONSUMER_DELAY_WHEN_MESSAGE_NOT_FOUND);
String tag = props.getProperty(RocketMQConfig.CONSUMER_TAG, RocketMQConfig.DEFAULT_CONSUMER_TAG);
int pullPoolSize = getInteger(props, RocketMQConfig.CONSUMER_PULL_POOL_SIZE,
RocketMQConfig.DEFAULT_CONSUMER_PULL_POOL_SIZE);
int pullBatchSize = getInteger(props, RocketMQConfig.CONSUMER_BATCH_SIZE,
RocketMQConfig.DEFAULT_CONSUMER_BATCH_SIZE);
pullConsumerScheduleService.setPullThreadNums(pullPoolSize);
pullConsumerScheduleService.registerPullTaskCallback(topic, new PullTaskCallback() {
@Override
public void doPullTask(MessageQueue mq, PullTaskContext pullTaskContext) {
try {
long offset = getMessageQueueOffset(mq);
if (offset < 0) {
return;
}
PullResult pullResult = consumer.pull(mq, tag, offset, pullBatchSize);
boolean found = false;
switch (pullResult.getPullStatus()) {
case FOUND:
List<MessageExt> messages = pullResult.getMsgFoundList();
for (MessageExt msg : messages) {
byte[] key = msg.getKeys() != null ? msg.getKeys().getBytes(StandardCharsets.UTF_8) : null;
byte[] value = msg.getBody();
OUT data = schema.deserializeKeyAndValue(key, value);
synchronized (lock) {
context.collectWithTimestamp(data, msg.getBornTimestamp());
}
}
found = true;
break;
case NO_MATCHED_MSG:
LOG.debug("No matched message after offset {} for queue {}", offset, mq);
break;
case NO_NEW_MSG:
break;
case OFFSET_ILLEGAL:
LOG.warn("Offset {} is illegal for queue {}", offset, mq);
break;
default:
break;
}
synchronized (lock) {
putMessageQueueOffset(mq, pullResult.getNextBeginOffset());
}
if (found) {
pullTaskContext.setPullNextDelayTimeMillis(0);
} else {
pullTaskContext.setPullNextDelayTimeMillis(delayWhenMessageNotFound);
}
} catch (Exception e) {
throw new RuntimeException(e);
}
}
});
try {
pullConsumerScheduleService.start();
} catch (MQClientException e) {
throw new RuntimeException(e);
}
runningChecker.setRunning(true);
awaitTermination();
}
private void awaitTermination() throws InterruptedException {
while (runningChecker.isRunning()) {
Thread.sleep(50);
}
}
private long getMessageQueueOffset(MessageQueue mq) throws MQClientException {
Long offset = offsetTable.get(mq);
if (restored && offset == null) {
offset = restoredOffsets.get(mq);
}
if (offset == null) {
offset = consumer.fetchConsumeOffset(mq, false);
if (offset < 0) {
String initialOffset = props.getProperty(RocketMQConfig.CONSUMER_OFFSET_RESET_TO, CONSUMER_OFFSET_LATEST);
switch (initialOffset) {
case CONSUMER_OFFSET_EARLIEST:
offset = consumer.minOffset(mq);
break;
case CONSUMER_OFFSET_LATEST:
offset = consumer.maxOffset(mq);
break;
case CONSUMER_OFFSET_TIMESTAMP:
offset = consumer.searchOffset(mq, getLong(props,
RocketMQConfig.CONSUMER_OFFSET_FROM_TIMESTAMP, System.currentTimeMillis()));
break;
default:
throw new IllegalArgumentException("Unknown value for CONSUMER_OFFSET_RESET_TO.");
}
}
}
offsetTable.put(mq, offset);
return offsetTable.get(mq);
}
private void putMessageQueueOffset(MessageQueue mq, long offset) throws MQClientException {
offsetTable.put(mq, offset);
if (!enableCheckpoint) {
consumer.updateConsumeOffset(mq, offset);
}
}
@Override
public void cancel() {
LOG.debug("cancel ...");
runningChecker.setRunning(false);
if (pullConsumerScheduleService != null) {
pullConsumerScheduleService.shutdown();
}
offsetTable.clear();
restoredOffsets.clear();
pendingOffsetsToCommit.clear();
}
@Override
public void close() throws Exception {
LOG.debug("close ...");
try {
cancel();
} finally {
super.close();
}
}
@Override
public void snapshotState(FunctionSnapshotContext context) throws Exception {
if (!runningChecker.isRunning()) {
LOG.debug("snapshotState() called on closed source; returning null.");
return;
}
if (LOG.isDebugEnabled()) {
LOG.debug("Snapshotting state {} ...", context.getCheckpointId());
}
unionOffsetStates.clear();
HashMap<MessageQueue, Long> currentOffsets = new HashMap<>(offsetTable.size());
Set<MessageQueue> assignedQueues = consumer.fetchMessageQueuesInBalance(topic);
offsetTable.entrySet().removeIf(item -> !assignedQueues.contains(item.getKey()));
for (Map.Entry<MessageQueue, Long> entry : offsetTable.entrySet()) {
unionOffsetStates.add(Tuple2.of(entry.getKey(), entry.getValue()));
currentOffsets.put(entry.getKey(), entry.getValue());
}
pendingOffsetsToCommit.put(context.getCheckpointId(), currentOffsets);
if (LOG.isDebugEnabled()) {
LOG.debug("Snapshotted state, last processed offsets: {}, checkpoint id: {}, timestamp: {}",
offsetTable, context.getCheckpointId(), context.getCheckpointTimestamp());
}
}
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
LOG.debug("initialize State ...");
this.unionOffsetStates = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<>(
OFFSETS_STATE_NAME, TypeInformation.of(new TypeHint<Tuple2<MessageQueue, Long>>() { })));
this.restored = context.isRestored();
if (restored) {
if (restoredOffsets == null) {
restoredOffsets = new ConcurrentHashMap<>();
}
for (Tuple2<MessageQueue, Long> mqOffsets : unionOffsetStates.get()) {
if (!restoredOffsets.containsKey(mqOffsets.f0) || restoredOffsets.get(mqOffsets.f0) < mqOffsets.f1) {
restoredOffsets.put(mqOffsets.f0, mqOffsets.f1);
}
}
LOG.info("Setting restore state in the consumer. Using the following offsets: {}", restoredOffsets);
} else {
LOG.info("No restore state for the consumer.");
}
}
@Override
public TypeInformation<OUT> getProducedType() {
return schema.getProducedType();
}
@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
if (!runningChecker.isRunning()) {
LOG.debug("notifyCheckpointComplete() called on closed source; returning null.");
return;
}
final int posInMap = pendingOffsetsToCommit.indexOf(checkpointId);
if (posInMap == -1) {
LOG.warn("Received confirmation for unknown checkpoint id {}", checkpointId);
return;
}
Map<MessageQueue, Long> offsets = (Map<MessageQueue, Long>)pendingOffsetsToCommit.remove(posInMap);
for (int i = 0; i < posInMap; i++) {
pendingOffsetsToCommit.remove(0);
}
if (offsets == null || offsets.size() == 0) {
LOG.debug("Checkpoint state was empty.");
return;
}
for (Map.Entry<MessageQueue, Long> entry : offsets.entrySet()) {
consumer.updateConsumeOffset(entry.getKey(), entry.getValue());
}
}
}SinkFunction
public interface SinkFunction<IN> extends Function, Serializable {
@Deprecated
default void invoke(IN value) throws Exception {}
default void invoke(IN value, Context context) throws Exception {
invoke(value);
}
@Public
interface Context<T> {
long currentProcessingTime();
long currentWatermark();
Long timestamp();
}
}@Public
public class DiscardingSink<T> implements SinkFunction<T> {
private static final long serialVersionUID = 1L;
@Override
public void invoke(T value) {}
}HBase Sink
public class CalendarHBaseSink extends RichSinkFunction<JSONObject> {
private static final long serialVersionUID = -6140896663267559061L;
private static final Logger LOGGER = LoggerFactory.getLogger(CalendarHBaseSink.class);
private static byte[] CF_BYTES = Bytes.toBytes("f");
private static TableName TABLE_NAME = TableName.valueOf("dw_hbase:calendar_record");
private Connection connection = null;
@Override
public void open(Configuration parameters) throws Exception {
super.open(parameters);
// HBase连接封装为单例
// RichSinkFunction.open()方法是按并行度执行的，而创建HBase连接是个很贵的操作
connection = HBaseConnection.get();
LOGGER.info("Opened CalendarHBaseSink");
}
@Override
public void close() throws Exception {
super.close();
HBaseConnection.close();
LOGGER.info("Closed CalendarHBaseSink");
}
// 数据预先解析为JSON
@Override
public void invoke(JSONObject record, Context context) throws Exception {
Integer eventType = record.getInteger("eventtype");
// 获取HBase中的列名映射，列名尽量短就是了
String qualifier = EventType.getColumnQualifier(eventType);
if (qualifier == null) {
return;
}
Long uid = record.getLong("uid");
Integer recordDate = record.getInteger("dateline");
String data = record.getString("data");
Long uploadTime = record.getLong("updatetime");
// try-with-resources语法。创建Table就很轻量级了
// 为了提高写入效率，在并发大时还可以使用HBase的BufferedMutator
try (Table table = connection.getTable(TABLE_NAME)) {
// 以UID和记录日期作为主键。注意设计一个好的RowKey
Put put = new Put(Bytes.toBytes(RowKeyUtil.getForCalendar(uid, recordDate)));
// 将上传时间作为HBase时间戳写入
put.addColumn(CF_BYTES, Bytes.toBytes(qualifier), uploadTime * 1000, Bytes.toBytes(data));
table.put(put);
}
}
}Streaming主程序
public static void main(String[] args) throws Exception {
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
env.enableCheckpointing(300000);
env.getCheckpointConfig().setCheckpointTimeout(60000);
// at least once的检查点就够用了
env.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.AT_LEAST_ONCE);
// 逻辑很简单，并且RocketMQ会记录offset，不必非要用FsStateBackend
env.setStateBackend(new MemoryStateBackend(true));
Properties consumerProps = new Properties();
consumerProps.setProperty(NAME_SERVER_ADDR, RocketMQConst.NAME_SERVER_TEST);
consumerProps.setProperty(CONSUMER_OFFSET_RESET_TO, "latest");
consumerProps.setProperty(CONSUMER_TOPIC, "calendar");
consumerProps.setProperty(CONSUMER_TAG, "*");
consumerProps.setProperty(CONSUMER_GROUP, "FLINK_STREAM_CALENDAR_TEST_1");
env.addSource(new RocketMQSource<>(new JSONDeserializationSchema(), consumerProps))
.name("calendar-rocketmq-source")
.addSink(new CalendarHBaseSink())
.name("calendar-hbase-sink");
env.execute();
} public static void main(String[] args) throws Exception {
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
env.enableCheckpointing(300000);
env.getCheckpointConfig().setCheckpointTimeout(60000);
// at least once的检查点就够用了
env.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.AT_LEAST_ONCE);
// 逻辑很简单，并且RocketMQ会记录offset，不必非要用FsStateBackend
env.setStateBackend(new MemoryStateBackend(true));
Properties consumerProps = new Properties();
consumerProps.setProperty(NAME_SERVER_ADDR, RocketMQConst.NAME_SERVER_TEST);
consumerProps.setProperty(CONSUMER_OFFSET_RESET_TO, "latest");
consumerProps.setProperty(CONSUMER_TOPIC, "calendar");
consumerProps.setProperty(CONSUMER_TAG, "*");
consumerProps.setProperty(CONSUMER_GROUP, "FLINK_STREAM_CALENDAR_TEST_1");
env.addSource(new RocketMQSource<>(new JSONDeserializationSchema(), consumerProps))
.name("calendar-rocketmq-source")
.addSink(new CalendarHBaseSink())
.name("calendar-hbase-sink");
env.execute();
}使用KeyedStream和分键状态（Keyed State），参见http://apache-flink-user-mailing-list-archive.2336050.n4.nabble.com/Maintaining-watermarks-per-key-instead-of-per-operator-instance-td7288.html。由于我们采用(UID, 日期)的双字段作为Key，状态空间有可能会奇大无比，目前持保留意见。
利用自带时间戳机制的外部存储。HBase的每个Cell都带有时间戳，在建表时设定版本数为1，就可以保证只留下最新的那条数据，这个对我们来说显然更实用。
文章不错？点个【在看】吧！ 👇