Flink-反压-4.源码分析-浮动缓冲区和专属缓冲区

在前面文章Flink-反压-2.源码分析-流程-1，我们知道BufferManager主要回收的是专属缓冲区，浮动缓冲区的递归回收，其实还是要看LocalBufferPool的逻辑

一.添加浮动缓冲区

1.BufferManager干了啥

其实就是调用BufferPool的requestBuffer()

/* 被RemoteInputChannel.onSenderBacklog()调用 用于请求浮动缓冲区
*    public void onSenderBacklog(int backlog) throws IOException {
*        notifyBufferAvailable(bufferManager.requestFloatingBuffers(backlog + initialCredit));
*     }
* */
int requestFloatingBuffers(int numRequired) {
    int numRequestedBuffers = 0;
    synchronized (bufferQueue) {
        // Similar to notifyBufferAvailable(), make sure that we never add a buffer after
        // channel
        // released all buffers via releaseAllResources().
        if (inputChannel.isReleased()) {
            return numRequestedBuffers;
        }

        numRequiredBuffers = numRequired;
        numRequestedBuffers = tryRequestBuffers();
    }
    return numRequestedBuffers;
}

private int tryRequestBuffers() {
    assert Thread.holdsLock(bufferQueue);

    int numRequestedBuffers = 0;
    while (bufferQueue.getAvailableBufferSize() < numRequiredBuffers
            && !isWaitingForFloatingBuffers) {
        BufferPool bufferPool = inputChannel.inputGate.getBufferPool();
        Buffer buffer = bufferPool.requestBuffer(); // 这是生产浮动缓冲区的核心逻辑
        if (buffer != null) {
            bufferQueue.addFloatingBuffer(buffer);
            numRequestedBuffers++;
        } else if (bufferPool.addBufferListener(this)) {
            isWaitingForFloatingBuffers = true;
            break;
        }
    }
    return numRequestedBuffers;
}

2.BufferPool是一个接口

public interface BufferPool extends BufferProvider, BufferRecycler{
    。。。
    其中，requestBuffer()是BufferProvider定义的
}

其实现类如下

(1) LocalBufferPool生产浮动缓冲区

@Override
public Buffer requestBuffer() {
    // 调toBuffer()
    return toBuffer(requestMemorySegment());
}

// 根据MemorySegment去生成一个NetworkBuffer，这个缓冲区就是刚才的浮动缓冲区
private Buffer toBuffer(MemorySegment memorySegment) {
    if (memorySegment == null) {
        return null;
    }
    return new NetworkBuffer(memorySegment, this);
}

二.释放各类缓冲区

1.释放浮动缓冲区

(1) BufferManager干了啥

由BufferManager的releaseFloatingBuffers()调用

逐个回收所有的浮动缓冲区，调用Buffer.recycleBuffer()，但其实调用的是LocalBufferPool.recycle()将浮动缓冲区回收

    // 释放浮动缓冲区，该方法被RemoteInputChannel.onBlockingUpstream()调用
    void releaseFloatingBuffers() {
        Queue<Buffer> buffers;
        synchronized (bufferQueue) {
            numRequiredBuffers = 0;
            buffers = bufferQueue.clearFloatingBuffers(); // 清空浮动缓冲区队列
        }

        // recycle all buffers out of the synchronization block to avoid dead lock
        // 逐个回收所有的浮动缓冲区，调用buffer.recycleBuffer()，但其实调用的是LocalBufferPool.recycle()将浮动缓冲区回收
        // 直接回收给本地缓存LocalBufferPool
        while (!buffers.isEmpty()) {
            buffers.poll().recycleBuffer();
        }
    }

(2) LocalBufferPool干了啥 ---TODO去补充其他的方法

总结：回收内存段时，优先分配给等待的监听器（上游生产者），分配成功则退出循环；若没有监听器或分配失败，则将内存段放入可用队列，并在系统从 "不可用" 转为 "可用" 时触发反压解除通知。

<1> recycle()

具体步骤如下：

1. 更新子分区缓冲区计数，针对channel已知的情况，进行if (subpartitionBuffersCount[channel]-- == maxBuffersPerChannel) { unavailableSubpartitionsCount--;}
2. 检查是否需要销毁或释放内存，需要，则调returnMemorySegment()去释放内存
3. 获取等待缓冲区的监视器
4. 当等待缓冲区的监听器为null时
   • 将浮动缓冲区加入可用队列
   • 当availabilityHelper之前处于不可用状态，且现在应该变为可用状态时，调availabilityHelper.getUnavailableToResetAvailable()去修改availableFuture为可用状态
5. 将可用的浮动缓冲区分配给下游，分配成功，则退出循环；否则，继续循环
6. 等待availableFuture状态切换完成

private void recycle(MemorySegment segment, int channel) {
    BufferListener listener;
    CompletableFuture<?> toNotify = null;
    do {
        // availableMemorySegments：可用内存段队列，存放已回收的MemorySegment
        synchronized (availableMemorySegments) {
            // 步1. 更新子分区缓冲区计数，针对channel已知的情况
            if (channel != UNKNOWN_CHANNEL) {
                // subpartitionBuffersCount：各子分区当前是由的缓冲区数量（用于限流）
                // maxBuffersPerChannel：单个子分区允许的最大缓冲区数量（反压阈值）。
                // 注意：下面的if条件是先判断再--，
                // 说明当回收一个缓冲区时，如果该channel的子分区缓冲区数量从maxBuffersPerChannel变为maxBuffersPerChannel-1了，说明有一个缓冲区可用了
                if (subpartitionBuffersCount[channel]-- == maxBuffersPerChannel) {
                    // unavailableSubpartitionsCount：当前不可用（缓冲区已满）的子分区数量。
                    unavailableSubpartitionsCount--;
                }
            }
            // 步2. 检查是否需要销毁或释放内存，需要，则调returnMemorySegment()去释放内存
            // hasExcessBuffers就是检查当前分配的缓冲区总数numberOfRequestedMemorySegments 是否> 池currentPoolSize的数量
            if (isDestroyed || hasExcessBuffers()) {
                returnMemorySegment(segment);
                return;
            } else {
                // 步3. 获取等待缓冲区的监视器
                // registeredListeners：等待缓冲区可用的监听器队列
                listener = registeredListeners.poll();
                // 步4.当等待缓冲区的监听器为null时
                if (listener == null) {
                    // 4.1 浮动缓冲区直接放入可用队列（无需更新子分区计数）
                    availableMemorySegments.add(segment);
                    //  4.2 当availabilityHelper之前处于不可用状态，且现在应该变为可用状态时，调availabilityHelper.getUnavailableToResetAvailable()去修改availableFuture为可用状态
                    // shouldBeAvailable==>!availableMemorySegments.isEmpty() && unavailableSubpartitionsCount == 0;
                    if (!availabilityHelper.isApproximatelyAvailable() && shouldBeAvailable()) {
                        // 将availableFuture由不可用调整为可用，并保存不可用状态给toNotify
                        toNotify = availabilityHelper.getUnavailableToResetAvailable();
                    }
                    break;
                }
            }

            checkConsistentAvailability();
        }
    } while (!fireBufferAvailableNotification(listener, segment)); // 步5. 循环条件
    // fireBufferAvailableNotification其实底层调的是BufferManager.notifyBufferAvailable()，就是将可用的浮动缓冲区分配给下游，分配成功，则退出循环；否则，继续循环
    // 步6. 等待状态切换完毕
    mayNotifyAvailable(toNotify);
}

<2> hasExcessBuffers()

// 持有availableMemorySegments锁，防止线程安全问题
@GuardedBy("availableMemorySegments")
private boolean hasExcessBuffers() {
    // 判断是否存在过剩缓冲区
    return numberOfRequestedMemorySegments > currentPoolSize;
}

<3> returnMemorySegment()

@GuardedBy("availableMemorySegments")
private void returnMemorySegment(MemorySegment segment) {
    assert Thread.holdsLock(availableMemorySegments);
    // 减少已申请的内存段计数
    numberOfRequestedMemorySegments--;
    // 将内存段归还给网络缓冲区池
    networkBufferPool.recyclePooledMemorySegment(segment);
}

<4> shouldBeAvailable()

@GuardedBy("availableMemorySegments")
private boolean shouldBeAvailable() {
    assert Thread.holdsLock(availableMemorySegments);
    // !availableMemorySegments.isEmpty()：表示存在可用的内存段集合
    // unavailableSubpartitionsCount == 0：表示所有子分区都已经可用
    return !availableMemorySegments.isEmpty() && unavailableSubpartitionsCount == 0;
}

<5> fireBufferAvailableNotification()

listener.notifyBufferAvailable()其实就是调的BufferManager.notifyBufferAvailable()

private boolean fireBufferAvailableNotification(
        BufferListener listener, MemorySegment segment) {
    // 其实底层调的是BufferManager.notifyBufferAvailable()
    return listener.notifyBufferAvailable(new NetworkBuffer(segment, this));
}

补充，BufferManager.notifyBufferAvailable()

@Override
public boolean notifyBufferAvailable(Buffer buffer) {
    if (inputChannel.isReleased()) {
        return false;
    }

    int numBuffers = 0;
    boolean isBufferUsed = false;
    try {
        synchronized (bufferQueue) {
            // 检查当前通道是否处于"等待浮动缓冲区"状态
            checkState(
                    isWaitingForFloatingBuffers,
                    "This channel should be waiting for floating buffers.");
            isWaitingForFloatingBuffers = false;// 重置等待状态

            // 再次检查通道状态（双重校验，避免并发释放），若可用的缓冲区数量 >= 需求缓冲区数量，直接return
            if (inputChannel.isReleased()
                    || bufferQueue.getAvailableBufferSize() >= numRequiredBuffers) {
                return false;
            }
            // 将可用的浮动缓冲区加入队列
            bufferQueue.addFloatingBuffer(buffer);
            isBufferUsed = true;// 标记缓冲区已被使用
            // 尝试请求更多缓冲区（按需扩展）
            numBuffers += 1 + tryRequestBuffers();
            bufferQueue.notifyAll(); // 唤醒等待缓冲区的线程（如下游算子取数据）
        }

        inputChannel.notifyBufferAvailable(numBuffers); // 调RemoteInputChannel.notifyBufferAvailable()通知上游，更新信用值
    } catch (Throwable t) {
        inputChannel.setError(t);
    }

    return isBufferUsed;// 返回是否成功使用缓冲区
}

<6> mayNotifyAvailable()

private void mayNotifyAvailable(@Nullable CompletableFuture<?> toNotify) {
    // 等待状态切换完成
    if (toNotify != null) {
        toNotify.complete(null);
    }
}

2.释放所有缓冲区

// 释放所有缓冲区（专属 + 浮动）
void releaseAllBuffers(ArrayDeque<Buffer> buffers) throws IOException {
    // Gather all exclusive buffers and recycle them to global pool in batch, because
    // we do not want to trigger redistribution of buffers after each recycle.
    final List<MemorySegment> exclusiveRecyclingSegments = new ArrayList<>();

    Exception err = null;
    Buffer buffer;
    // 1.回收参数传入的 `buffers` 队列
    while ((buffer = buffers.poll()) != null) {
        try {
            if (buffer.getRecycler() == BufferManager.this) {
                // 专属缓冲区：收集内存段，后续批量归还给全局池
                exclusiveRecyclingSegments.add(buffer.getMemorySegment());
            } else {
                // 浮动缓冲区：直接回收，调LocalBufferPool.recycleBuffer()
                buffer.recycleBuffer();
            }
        } catch (Exception e) {
            err = firstOrSuppressed(e, err);
        }
    }
    // 2.回收`bufferQueue` 内部的缓冲区
    try {
        synchronized (bufferQueue) {
            bufferQueue.releaseAll(exclusiveRecyclingSegments);
            bufferQueue.notifyAll();
        }
    } catch (Exception e) {
        err = firstOrSuppressed(e, err);
    }
    try {
        // 批量把专属缓冲区归还给全局池
        if (exclusiveRecyclingSegments.size() > 0) {
            globalPool.recycleUnpooledMemorySegments(exclusiveRecyclingSegments);
        }
    } catch (Exception e) {
        err = firstOrSuppressed(e, err);
    }
    if (err != null) {
        throw err instanceof IOException ? (IOException) err : new IOException(err);
    }
}
    
// 调用的内部类的releaseAll()
void releaseAll(List<MemorySegment> exclusiveSegments) {
    Buffer buffer;
    while ((buffer = floatingBuffers.poll()) != null) {
        // 浮动缓冲区调LocalBufferPool.recycleBuffer()回收
        buffer.recycleBuffer();
    }
    while ((buffer = exclusiveBuffers.poll()) != null) {
        // 专属缓冲区将内存段加入exclusiveSegments
        exclusiveSegments.add(buffer.getMemorySegment());
    }
}