<commit_before>/*
 * Copyright 2004-present Facebook. All Rights Reserved.
 */
package com.facebook.presto.execution;

import com.google.common.base.Objects;
import com.google.common.base.Preconditions;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.primitives.Ints;
import io.airlift.units.Duration;

import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.NotThreadSafe;
import javax.annotation.concurrent.ThreadSafe;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import static com.facebook.presto.execution.BufferResult.emptyResults;
import static java.util.concurrent.TimeUnit.NANOSECONDS;

@ThreadSafe
public class SharedBuffer<T>
{
    public static enum QueueState
    {
        /**
         * Additional queues can be added.
         */
        OPEN,
        /**
         * No more queues can be added.
         */
        NO_MORE_QUEUES,
        /**
         * No more queues can be added and all elements have been consumed.
         */
        FINISHED
    }

    @GuardedBy("this")
    private final LinkedList<T> masterQueue;
    @GuardedBy("this")
    private long masterSequenceId;
    @GuardedBy("this")
    private Map<String, NamedQueue> namedQueues = new HashMap<>();
    @GuardedBy("this")
    private final SortedSet<NamedQueue> openQueuesBySequenceId = new TreeSet<>();
    @GuardedBy("this")
    private QueueState state = QueueState.OPEN;

    private final AtomicLong pagesAdded = new AtomicLong();

    /**
     * If true, no more elements can be added to the queue.
     */
    private final AtomicBoolean closed = new AtomicBoolean();

    private final Semaphore notFull;

    public SharedBuffer(int capacity)
    {
        Preconditions.checkArgument(capacity > 0, "Capacity must be at least 1");
        masterQueue = new LinkedList<>();
        notFull = new Semaphore(capacity);
    }

    public synchronized boolean isFinished()
    {
        return state == QueueState.FINISHED;
    }

    public synchronized SharedBufferInfo getInfo()
    {
        ImmutableList.Builder<BufferInfo> infos = ImmutableList.builder();
        for (NamedQueue namedQueue : namedQueues.values()) {
            infos.add(new BufferInfo(namedQueue.getQueueId(), namedQueue.isFinished(), namedQueue.size(), namedQueue.elementsRemoved()));
        }
        return new SharedBufferInfo(state, pagesAdded.get(), infos.build());
    }

    public synchronized void addQueue(String queueId)
    {
        Preconditions.checkNotNull(queueId, "queueId is null");
        if (namedQueues.containsKey(queueId)) {
            return;
        }
        Preconditions.checkState(state == QueueState.OPEN, "%s is not OPEN", SharedBuffer.class.getSimpleName());
        NamedQueue namedQueue = new NamedQueue(queueId);
        namedQueues.put(queueId, namedQueue);<commit_msg>Ignore add output buffer calls when query is finished<commit_after>/*
 * Copyright 2004-present Facebook. All Rights Reserved.
 */
package com.facebook.presto.execution;

import com.google.common.base.Objects;
import com.google.common.base.Preconditions;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.primitives.Ints;
import io.airlift.units.Duration;

import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.NotThreadSafe;
import javax.annotation.concurrent.ThreadSafe;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import static com.facebook.presto.execution.BufferResult.emptyResults;
import static java.util.concurrent.TimeUnit.NANOSECONDS;

@ThreadSafe
public class SharedBuffer<T>
{
    public static enum QueueState
    {
        /**
         * Additional queues can be added.
         */
        OPEN,
        /**
         * No more queues can be added.
         */
        NO_MORE_QUEUES,
        /**
         * No more queues can be added and all elements have been consumed.
         */
        FINISHED
    }

    @GuardedBy("this")
    private final LinkedList<T> masterQueue;
    @GuardedBy("this")
    private long masterSequenceId;
    @GuardedBy("this")
    private Map<String, NamedQueue> namedQueues = new HashMap<>();
    @GuardedBy("this")
    private final SortedSet<NamedQueue> openQueuesBySequenceId = new TreeSet<>();
    @GuardedBy("this")
    private QueueState state = QueueState.OPEN;

    private final AtomicLong pagesAdded = new AtomicLong();

    /**
     * If true, no more elements can be added to the queue.
     */
    private final AtomicBoolean closed = new AtomicBoolean();

    private final Semaphore notFull;

    public SharedBuffer(int capacity)
    {
        Preconditions.checkArgument(capacity > 0, "Capacity must be at least 1");
        masterQueue = new LinkedList<>();
        notFull = new Semaphore(capacity);
    }

    public synchronized boolean isFinished()
    {
        return state == QueueState.FINISHED;
    }

    public synchronized SharedBufferInfo getInfo()
    {
        ImmutableList.Builder<BufferInfo> infos = ImmutableList.builder();
        for (NamedQueue namedQueue : namedQueues.values()) {
            infos.add(new BufferInfo(namedQueue.getQueueId(), namedQueue.isFinished(), namedQueue.size(), namedQueue.elementsRemoved()));
        }
        return new SharedBufferInfo(state, pagesAdded.get(), infos.build());
    }

    public synchronized void addQueue(String queueId)
    {
        Preconditions.checkNotNull(queueId, "queueId is null");
        // ignore buffers added after query finishes, which can happen when a query is canceled
        // also ignore duplicates, which is normal
        if (state == QueueState.FINISHED || namedQueues.containsKey(queueId)) {
            return;
        }
        Preconditions.checkState(state == QueueState.OPEN, "%s is not OPEN", SharedBuffer.class.getSimpleName());
        NamedQueue namedQueue = new NamedQueue(queueId);
        namedQueues.put(queueId, namedQueue);
{'fraction_non_alphanumeric': 0.10058287251531908, 'fraction_numerical': 0.0017934538932894933, 'mean_word_length': 3.3398184176394294, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '377525', 'n_tokens_mistral': 1915, 'n_tokens_neox': 1845, 'n_words': 461}