All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.neo4j.driver.internal.handlers.LegacyPullAllResponseHandler Maven / Gradle / Ivy

There is a newer version: 5.27.0
Show newest version
/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [https://neo4j.com]
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.neo4j.driver.internal.handlers;

import static java.util.Collections.emptyMap;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.CompletableFuture.completedFuture;
import static org.neo4j.driver.internal.util.Futures.completedWithNull;
import static org.neo4j.driver.internal.util.Futures.failedFuture;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
import org.neo4j.driver.Query;
import org.neo4j.driver.Record;
import org.neo4j.driver.Value;
import org.neo4j.driver.exceptions.Neo4jException;
import org.neo4j.driver.internal.InternalRecord;
import org.neo4j.driver.internal.messaging.request.PullAllMessage;
import org.neo4j.driver.internal.messaging.v55.BoltProtocolV55;
import org.neo4j.driver.internal.spi.Connection;
import org.neo4j.driver.internal.util.Futures;
import org.neo4j.driver.internal.util.Iterables;
import org.neo4j.driver.internal.util.MetadataExtractor;
import org.neo4j.driver.summary.GqlStatusObject;
import org.neo4j.driver.summary.ResultSummary;

/**
 * This is the Pull All response handler that handles pull all messages in Bolt v3 and previous protocol versions.
 */
public class LegacyPullAllResponseHandler extends AbstractRecordStateResponseHandler implements PullAllResponseHandler {
    private static final Queue UNINITIALIZED_RECORDS = Iterables.emptyQueue();

    static final int RECORD_BUFFER_LOW_WATERMARK = Integer.getInteger("recordBufferLowWatermark", 300);
    static final int RECORD_BUFFER_HIGH_WATERMARK = Integer.getInteger("recordBufferHighWatermark", 1000);

    private final Query query;
    private final RunResponseHandler runResponseHandler;
    protected final MetadataExtractor metadataExtractor;
    protected final Connection connection;
    private final PullResponseCompletionListener completionListener;

    // initialized lazily when first record arrives
    private Queue records = UNINITIALIZED_RECORDS;

    private boolean autoReadManagementEnabled = true;
    private boolean finished;
    private Throwable failure;
    private ResultSummary summary;

    private boolean ignoreRecords;
    private CompletableFuture recordFuture;
    private CompletableFuture failureFuture;

    public LegacyPullAllResponseHandler(
            Query query,
            RunResponseHandler runResponseHandler,
            Connection connection,
            MetadataExtractor metadataExtractor,
            PullResponseCompletionListener completionListener) {
        this.query = requireNonNull(query);
        this.runResponseHandler = requireNonNull(runResponseHandler);
        this.metadataExtractor = requireNonNull(metadataExtractor);
        this.connection = requireNonNull(connection);
        this.completionListener = requireNonNull(completionListener);
    }

    @Override
    public boolean canManageAutoRead() {
        return true;
    }

    @Override
    public synchronized void onSuccess(Map metadata) {
        finished = true;
        Neo4jException exception = null;
        try {
            summary = extractResultSummary(
                    metadata,
                    generateGqlStatusObject(runResponseHandler.queryKeys().keys()));
        } catch (Neo4jException e) {
            exception = e;
        }

        if (exception == null) {
            completionListener.afterSuccess(metadata);

            completeRecordFuture(null);
            completeFailureFuture(null);
        } else {
            onFailure(exception);
        }
    }

    @Override
    public synchronized void onFailure(Throwable error) {
        finished = true;
        summary = extractResultSummary(emptyMap(), null);

        completionListener.afterFailure(error);

        var failedRecordFuture = failRecordFuture(error);
        if (failedRecordFuture) {
            // error propagated through the record future
            completeFailureFuture(null);
        } else {
            var completedFailureFuture = completeFailureFuture(error);
            if (!completedFailureFuture) {
                // error has not been propagated to the user, remember it
                failure = error;
            }
        }
    }

    @Override
    public synchronized void onRecord(Value[] fields) {
        recordState = RecordState.HAD_RECORD;
        if (ignoreRecords) {
            completeRecordFuture(null);
        } else {
            Record record = new InternalRecord(runResponseHandler.queryKeys(), fields);
            enqueueRecord(record);
            completeRecordFuture(record);
        }
    }

    @Override
    public synchronized void disableAutoReadManagement() {
        autoReadManagementEnabled = false;
    }

    public synchronized CompletionStage peekAsync() {
        var record = records.peek();
        if (record == null) {
            if (failure != null) {
                return failedFuture(extractFailure());
            }

            if (ignoreRecords || finished) {
                return completedWithNull();
            }

            if (recordFuture == null) {
                recordFuture = new CompletableFuture<>();
            }
            return recordFuture;
        } else {
            return completedFuture(record);
        }
    }

    public synchronized CompletionStage nextAsync() {
        return peekAsync().thenApply(ignore -> dequeueRecord());
    }

    public synchronized CompletionStage consumeAsync() {
        ignoreRecords = true;
        records.clear();
        return pullAllFailureAsync().thenApply(error -> {
            if (error != null) {
                throw Futures.asCompletionException(error);
            }
            return summary;
        });
    }

    public synchronized  CompletionStage> listAsync(Function mapFunction) {
        return pullAllFailureAsync().thenApply(error -> {
            if (error != null) {
                throw Futures.asCompletionException(error);
            }
            return recordsAsList(mapFunction);
        });
    }

    @Override
    public void prePopulateRecords() {
        synchronized (this) {
            recordState = RecordState.NO_RECORD;
        }
        connection.writeAndFlush(PullAllMessage.PULL_ALL, this);
    }

    public synchronized CompletionStage pullAllFailureAsync() {
        if (failure != null) {
            return completedFuture(extractFailure());
        } else if (finished) {
            return completedWithNull();
        } else {
            if (failureFuture == null) {
                // neither SUCCESS nor FAILURE message has arrived, register future to be notified when it arrives
                // future will be completed with null on SUCCESS and completed with Throwable on FAILURE
                // enable auto-read, otherwise we might not read SUCCESS/FAILURE if records are not consumed
                enableAutoRead();
                failureFuture = new CompletableFuture<>();
            }
            return failureFuture;
        }
    }

    private void enqueueRecord(Record record) {
        if (records == UNINITIALIZED_RECORDS) {
            records = new ArrayDeque<>();
        }

        records.add(record);

        var shouldBufferAllRecords = failureFuture != null;
        // when failure is requested we have to buffer all remaining records and then return the error
        // do not disable auto-read in this case, otherwise records will not be consumed and trailing
        // SUCCESS or FAILURE message will not arrive as well, so callers will get stuck waiting for the error
        if (!shouldBufferAllRecords && records.size() > RECORD_BUFFER_HIGH_WATERMARK) {
            // more than high watermark records are already queued, tell connection to stop auto-reading from network
            // this is needed to deal with slow consumers, we do not want to buffer all records in memory if they are
            // fetched from network faster than consumed
            disableAutoRead();
        }
    }

    private Record dequeueRecord() {
        var record = records.poll();

        if (records.size() < RECORD_BUFFER_LOW_WATERMARK) {
            // less than low watermark records are now available in the buffer, tell connection to pre-fetch more
            // and populate queue with new records from network
            enableAutoRead();
        }

        return record;
    }

    private  List recordsAsList(Function mapFunction) {
        if (!finished) {
            throw new IllegalStateException("Can't get records as list because SUCCESS or FAILURE did not arrive");
        }

        List result = new ArrayList<>(records.size());
        while (!records.isEmpty()) {
            var record = records.poll();
            result.add(mapFunction.apply(record));
        }
        return result;
    }

    private Throwable extractFailure() {
        if (failure == null) {
            throw new IllegalStateException("Can't extract failure because it does not exist");
        }

        var error = failure;
        failure = null; // propagate failure only once
        return error;
    }

    private void completeRecordFuture(Record record) {
        if (recordFuture != null) {
            var future = recordFuture;
            recordFuture = null;
            future.complete(record);
        }
    }

    private boolean failRecordFuture(Throwable error) {
        if (recordFuture != null) {
            var future = recordFuture;
            recordFuture = null;
            future.completeExceptionally(error);
            return true;
        }
        return false;
    }

    private boolean completeFailureFuture(Throwable error) {
        if (failureFuture != null) {
            var future = failureFuture;
            failureFuture = null;
            future.complete(error);
            return true;
        }
        return false;
    }

    private ResultSummary extractResultSummary(Map metadata, GqlStatusObject gqlStatusObject) {
        var resultAvailableAfter = runResponseHandler.resultAvailableAfter();
        return metadataExtractor.extractSummary(
                query,
                connection,
                resultAvailableAfter,
                metadata,
                connection.protocol().version().compareTo(BoltProtocolV55.VERSION) < 0,
                gqlStatusObject);
    }

    private void enableAutoRead() {
        if (autoReadManagementEnabled) {
            connection.enableAutoRead();
        }
    }

    private void disableAutoRead() {
        if (autoReadManagementEnabled) {
            connection.disableAutoRead();
        }
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy