Optimizing Backpressure in Flink’s Network Stack

I often heard praise for Flink’s backpressure mechanism. Only after encountering a problem in an actual application and investigating it closely did I realize it was not perfect.

Network Stack

Several posts explain Apache Flink’s network stack:

• A Deep-Dive into Flink Network Stack
• Flink Principles and Implementation: Handling Backpressure

On the receiving side, a RemoteChannel has m+n buffers: floating plus exclusive buffers. Because each RemoteInputChannel has its own exclusive buffers, different channels in the same TaskManager do not block each other when receiving data.

From the receiving side, this looks sound: slow input channels do not block faster ones. From the sending side, however, things differ. For clarity, the following diagram omits the receiver’s credit-based mechanism:

Different subpartitions of the same task share a LocalBufferPool without a comparable credit-based mechanism. When sending data, they request a MemorySegment from that pool, wrap the data, and send it. Could a slowly draining subpartition occupy all pool resources, blocking a faster one that cannot acquire a buffer?

Yes. Try the demo below. I deliberately slow one task’s consumption. Since MapTask consumes more slowly than SourceTask produces, and distribution is round-robin, the subpartition destined for the slow MapTask occupies all LocalBufferPool resources. The subpartition destined for the fast MapTask cannot acquire resources and blocks.

1
@Test
2
def testRebalance(): Unit = {
3
  val env = StreamExecutionEnvironment.getExecutionEnvironment
4

5
  env.setParallelism(2)
6
  env.getConfig.setMaxParallelism(2)
7

8
  val sourceStream = env.addSource(new SourceFunction[(Int, Int)] {
9

10
    override def run(ctx: SourceContext[(Int, Int)]): Unit = {
11
      while (true) {
12
        0 until 10 foreach {
13
          // keys '1' and '2' hash to different buckets
14
          i => ctx.collect((1 + MathUtils.murmurHash(i) % 5, i))
15
        }
16
      }
17
    }
18

19
    override def cancel(): Unit = {}
20
  })
21

22
  sourceStream.rebalance.map(new RichMapFunction[(Int, Int), String] {
23

24
    var block = false
25

26
    override def open(parameters: Configuration): Unit = {
27
      super.open(parameters)
28
      if (getRuntimeContext.getIndexOfThisSubtask == 0) {
29
        block = true
30
      }
31
    }
32

33
    override def map(value: (Int, Int)): String = {
34
      if (block) {
35
        Thread.sleep(10000)
36
      }
37
      s"Subtask Index: ${getRuntimeContext.getIndexOfThisSubtask}"
38
    }
39
  }).print()
40

41
  println(env.getExecutionPlan)
42
  env.execute()
43
}

Optimization

The demo uses round-robin distribution, so I do not care which subpartition receives a record, only that later processing avoids skew and backpressure. This case can be optimized. Credit-based flow control tracks a sender backlog: data waiting to be sent. Receivers use backlog size to decide whether to request floating buffers. We can also use it to determine how many LocalBufferPool resources each subpartition occupies. Here is the key code:

Introduce a BacklogBasedSelector trait and implement it in RebalancePartitioner. RecordWriter checks whether its channelSelector implements the trait and, if so, whether targetChannel is blocked.

1
if (targetPartition instanceof ResultPartition && channelSelector instanceof BacklogBasedSelector) {
2
  targetChannel = ((ResultPartition) targetPartition).changIfBlockingChannel(targetChannel);
3
}

Add changeIfBlockingChannel to ResultPartition.

1
public int changIfBlockingChannel(int subpartitionIndex) {
2
  int bufferSize = bufferPool.getNumBuffers();
3
  int backlogSize = subpartitions[subpartitionIndex].getBuffersInBacklog();
4
  // 根据 bufferPool 的大小和 subpartition 的 backlog 大小来选择是否需要更换 targetChannel
5
}

This only works when receiving RemoteInputChannels are interchangeable. For operations such as keyBy, there is currently no better alternative here.