org.elasticsearch.compute.operator.exchange.ExchangeBuffer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of x-pack-esql-compute Show documentation
Elasticsearch subproject :x-pack:plugin:esql:compute
There is a newer version: 8.16.1
/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.compute.operator.exchange;

import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.support.SubscribableListener;
import org.elasticsearch.compute.data.Page;
import org.elasticsearch.compute.operator.Operator;

import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

final class ExchangeBuffer {

    private final Queue queue = new ConcurrentLinkedQueue<>();
    // uses a separate counter for size for CAS; and ConcurrentLinkedQueue#size is not a constant time operation.
    private final AtomicInteger queueSize = new AtomicInteger();
    private final int maxSize;

    private final Object notEmptyLock = new Object();
    private SubscribableListener notEmptyFuture = null;

    private final Object notFullLock = new Object();
    private SubscribableListener notFullFuture = null;

    private final SubscribableListener completionFuture = new SubscribableListener<>();

    private volatile boolean noMoreInputs = false;

    ExchangeBuffer(int maxSize) {
        if (maxSize < 1) {
            throw new IllegalArgumentException("max_buffer_size must be at least one; got=" + maxSize);
        }
        this.maxSize = maxSize;
    }

    void addPage(Page page) {
        queue.add(page);
        if (queueSize.incrementAndGet() == 1) {
            notifyNotEmpty();
        }
        if (noMoreInputs) {
            discardPages();
        }
    }

    Page pollPage() {
        final var page = queue.poll();
        if (page != null && queueSize.decrementAndGet() == maxSize - 1) {
            notifyNotFull();
        }
        if (page == null && noMoreInputs && queueSize.get() == 0) {
            completionFuture.onResponse(null);
        }
        return page;
    }

    private void notifyNotEmpty() {
        final SubscribableListener toNotify;
        synchronized (notEmptyLock) {
            toNotify = notEmptyFuture;
            notEmptyFuture = null;
        }
        if (toNotify != null) {
            toNotify.onResponse(null);
        }
    }

    private void notifyNotFull() {
        final SubscribableListener toNotify;
        synchronized (notFullLock) {
            toNotify = notFullFuture;
            notFullFuture = null;
        }
        if (toNotify != null) {
            toNotify.onResponse(null);
        }
    }

    SubscribableListener waitForWriting() {
        // maxBufferSize check is not water-tight as more than one sink can pass this check at the same time.
        if (queueSize.get() < maxSize || noMoreInputs) {
            return Operator.NOT_BLOCKED;
        }
        synchronized (notFullLock) {
            if (queueSize.get() < maxSize || noMoreInputs) {
                return Operator.NOT_BLOCKED;
            }
            if (notFullFuture == null) {
                notFullFuture = new SubscribableListener<>();
            }
            return notFullFuture;
        }
    }

    SubscribableListener waitForReading() {
        if (size() > 0 || noMoreInputs) {
            return Operator.NOT_BLOCKED;
        }
        synchronized (notEmptyLock) {
            if (size() > 0 || noMoreInputs) {
                return Operator.NOT_BLOCKED;
            }
            if (notEmptyFuture == null) {
                notEmptyFuture = new SubscribableListener<>();
            }
            return notEmptyFuture;
        }
    }

    private void discardPages() {
        Page p;
        while ((p = pollPage()) != null) {
            p.releaseBlocks();
        }
    }

    void finish(boolean drainingPages) {
        noMoreInputs = true;
        if (drainingPages) {
            discardPages();
        }
        notifyNotEmpty();
        if (drainingPages || queueSize.get() == 0) {
            completionFuture.onResponse(null);
        }
    }

    boolean isFinished() {
        return completionFuture.isDone();
    }

    boolean noMoreInputs() {
        return noMoreInputs;
    }

    int size() {
        return queueSize.get();
    }

    /**
     * Adds a listener that will be notified when this exchange buffer is finished.
     */
    void addCompletionListener(ActionListener listener) {
        completionFuture.addListener(listener);
    }
}