be/src/io/fs/stream_load_pipe.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

#pragma once

#include <gen_cpp/internal_service.pb.h>

#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <memory>
#include <mutex>
#include <string>

#include "common/status.h"
#include "core/custom_allocator.h"
#include "io/fs/file_reader.h"
#include "io/fs/path.h"
#include "load/message_body_sink.h"
#include "util/byte_buffer.h"
#include "util/slice.h"

namespace doris::io {
struct IOContext;

static inline constexpr size_t kMaxPipeBufferedBytes = 4 * 1024 * 1024;

class StreamLoadPipe : public MessageBodySink, public FileReader {
public:
    StreamLoadPipe(size_t max_buffered_bytes = kMaxPipeBufferedBytes,
                   size_t min_chunk_size = 64 * 1024, int64_t total_length = -1,
                   bool use_proto = false);
    ~StreamLoadPipe() override;

    Status append_and_flush(const char* data, size_t size, size_t proto_byte_size = 0);

    Status append(std::unique_ptr<PDataRow>&& row);
    Status append(const char* data, size_t size) override;
    Status append(const ByteBufferPtr& buf) override;

    virtual Status append_with_line_delimiter(const char* data, size_t size) {
        RETURN_IF_ERROR(append(data, size));
        return append("\n", 1);
    }

    virtual Status append_json(const char* data, size_t size) {
        return append_and_flush(data, size);
    }

    const Path& path() const override { return _path; }

    size_t size() const override { return 0; }

    int64_t mtime() const override { return 0; }

    // called when consumer finished
    Status close() override {
        if (!(_finished || _cancelled)) {
            cancel("closed");
        }
        return Status::OK();
    }

    bool closed() const override { return _cancelled; }

    // called when producer finished
    virtual Status finish() override;

    // called when producer/consumer failed
    virtual void cancel(const std::string& reason) override;

    Status read_one_message(DorisUniqueBufferPtr<uint8_t>* data, size_t* length);

    size_t get_queue_size() { return _buf_queue.size(); }

    // used for pipeline load, which use TUniqueId(lo: query_id.lo + fragment_id, hi: query_id.hi) as pipe_id
    static TUniqueId calculate_pipe_id(const UniqueId& query_id, int32_t fragment_id);

    size_t max_capacity() const { return _max_buffered_bytes; }

    size_t current_capacity();

    bool is_chunked_transfer() const { return _is_chunked_transfer; }

    void set_is_chunked_transfer(bool is_chunked_transfer) {
        _is_chunked_transfer = is_chunked_transfer;
    }

protected:
    Status read_at_impl(size_t offset, Slice result, size_t* bytes_read,
                        const IOContext* io_ctx) override;

private:
    // read the next buffer from _buf_queue
    Status _read_next_buffer(DorisUniqueBufferPtr<uint8_t>* data, size_t* length);

    Status _append(const ByteBufferPtr& buf, size_t proto_byte_size = 0);

    // Blocking queue
    std::mutex _lock;
    size_t _buffered_bytes;
    size_t _proto_buffered_bytes;
    size_t _max_buffered_bytes;
    size_t _min_chunk_size;
    // The total amount of data expected to be read.
    // In some scenarios, such as loading json format data through stream load,
    // the data needs to be completely read before it can be parsed,
    // so the total size of the data needs to be known.
    // The default is -1, which means that the data arrives in a stream
    // and the length is unknown.
    // size_t is unsigned, so use int64_t
    int64_t _total_length = -1;
    bool _use_proto = false;
    std::deque<ByteBufferPtr> _buf_queue;
    std::deque<std::unique_ptr<PDataRow>> _data_row_ptrs;
    std::condition_variable _put_cond;
    std::condition_variable _get_cond;

    ByteBufferPtr _write_buf;

    // no use, only for compatibility with the `Path` interface
    Path _path = "";

    // When importing JSON data and using chunked transfer encoding,
    // the data needs to be completely read before it can be parsed.
    bool _is_chunked_transfer = false;
};
} // namespace doris::io