// 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.

#include "util/thrift_server.h"

#include <glog/logging.h>
#include <thrift/Thrift.h>
#include <thrift/concurrency/ThreadFactory.h>
#include <thrift/concurrency/ThreadManager.h>
#include <thrift/protocol/TBinaryProtocol.h>
#include <thrift/protocol/TProtocol.h>
#include <thrift/server/TNonblockingServer.h>
#include <thrift/server/TThreadPoolServer.h>
#include <thrift/server/TThreadedServer.h>
#include <thrift/transport/TBufferTransports.h>
#include <thrift/transport/TNonblockingServerSocket.h>
#include <thrift/transport/TServerSocket.h>
#include <thrift/transport/TSocket.h>
#include <thrift/transport/TTransport.h>
// IWYU pragma: no_include <bits/chrono.h>
#include <chrono> // IWYU pragma: keep
#include <condition_variable>
#include <memory>
#include <mutex>
#include <sstream>
#include <thread>

#include "common/config.h"
#include "common/metrics/doris_metrics.h"
#include "service/backend_options.h"

namespace apache {
namespace thrift {
class TProcessor;

namespace transport {
class TServerTransport;
} // namespace transport
} // namespace thrift
} // namespace apache

namespace doris {

DEFINE_GAUGE_METRIC_PROTOTYPE_3ARG(thrift_current_connections, MetricUnit::CONNECTIONS,
                                   "Number of currently active connections");
DEFINE_COUNTER_METRIC_PROTOTYPE_3ARG(thrift_connections_total, MetricUnit::CONNECTIONS,
                                     "Total connections made over the lifetime of this server");

// Nonblocking Server socket implementation of TNonblockingServerTransport. Wrapper around a unix
// socket listen and accept calls.
class ImprovedNonblockingServerSocket : public apache::thrift::transport::TNonblockingServerSocket {
    using TConfiguration = apache::thrift::TConfiguration;
    using TSocket = apache::thrift::transport::TSocket;

public:
    // Constructor.
    ImprovedNonblockingServerSocket(int port)
            : TNonblockingServerSocket(port),
              config(std::make_shared<TConfiguration>(config::thrift_max_message_size)) {}
    ~ImprovedNonblockingServerSocket() override = default;

protected:
    std::shared_ptr<TSocket> createSocket(THRIFT_SOCKET clientSocket) override {
        return std::make_shared<TSocket>(clientSocket, config);
    }

private:
    std::shared_ptr<TConfiguration> config;
};

class ImprovedBufferedTransportFactory
        : public apache::thrift::transport::TBufferedTransportFactory {
    using TConfiguration = apache::thrift::TConfiguration;
    using TTransport = apache::thrift::transport::TTransport;
    using TBufferedTransport = apache::thrift::transport::TBufferedTransport;

public:
    ImprovedBufferedTransportFactory()
            : config(std::make_shared<TConfiguration>(config::thrift_max_message_size)) {}
    ~ImprovedBufferedTransportFactory() override = default;

    std::shared_ptr<TTransport> getTransport(std::shared_ptr<TTransport> trans) override {
        return std::make_shared<TBufferedTransport>(std::move(trans), config);
    }

private:
    std::shared_ptr<TConfiguration> config;
};

// Helper class that starts a server in a separate thread, and handles
// the inter-thread communication to monitor whether it started
// correctly.
class ThriftServer::ThriftServerEventProcessor
        : public apache::thrift::server::TServerEventHandler {
public:
    ThriftServerEventProcessor(ThriftServer* thrift_server)
            : _thrift_server(thrift_server), _signal_fired(false) {}

    // friendly to code style
    ~ThriftServerEventProcessor() override = default;

    // Called by TNonBlockingServer when server has acquired its resources and is ready to
    // serve, and signals to StartAndWaitForServer that start-up is finished.
    // From TServerEventHandler.
    void preServe() override;

    // Called when a client connects; we create per-client state and call any
    // SessionHandlerIf handler.
    void* createContext(std::shared_ptr<apache::thrift::protocol::TProtocol> input,
                        std::shared_ptr<apache::thrift::protocol::TProtocol> output) override;

    // Called when a client starts an RPC; we set the thread-local session key.
    void processContext(void* context,
                        std::shared_ptr<apache::thrift::transport::TTransport> output) override;

    // Called when a client disconnects; we call any SessionHandlerIf handler.
    void deleteContext(void* serverContext,
                       std::shared_ptr<apache::thrift::protocol::TProtocol> input,
                       std::shared_ptr<apache::thrift::protocol::TProtocol> output) override;

    // Waits for a timeout of TIMEOUT_MS for a server to signal that it has started
    // correctly.
    Status start_and_wait_for_server();

private:
    // Lock used to ensure that there are no missed notifications between starting the
    // supervision thread and calling _signal_cond.timed_wait. Also used to ensure
    // thread-safe access to members of _thrift_server
    std::mutex _signal_lock;

    // Condition variable that is notified by the supervision thread once either
    // a) all is well or b) an error occurred.
    std::condition_variable _signal_cond;

    // The ThriftServer under management. This class is a friend of ThriftServer, and
    // reaches in to change member variables at will.
    ThriftServer* _thrift_server;

    // Guards against spurious condition variable wakeups
    bool _signal_fired;

    // The time, in milliseconds, to wait for a server to come up
    const static int TIMEOUT_MS;

    // Called in a separate thread; wraps TNonBlockingServer::serve in an exception handler
    void supervise();
};

// https://stackoverflow.com/questions/5391973/undefined-reference-to-static-const-int
const int ThriftServer::ThriftServerEventProcessor::TIMEOUT_MS = 2500;

Status ThriftServer::ThriftServerEventProcessor::start_and_wait_for_server() {
    // Locking here protects against missed notifications if Supervise executes quickly
    std::unique_lock<std::mutex> lock(_signal_lock);
    _thrift_server->_started = false;

    _thrift_server->_server_thread = std::make_unique<std::thread>(
            &ThriftServer::ThriftServerEventProcessor::supervise, this);

    // Loop protects against spurious wakeup. Locks provide necessary fences to ensure
    // visibility.
    while (!_signal_fired) {
        // Yields lock and allows supervision thread to continue and signal
        std::cv_status cvsts = _signal_cond.wait_for(lock, std::chrono::milliseconds(TIMEOUT_MS));
        if (cvsts == std::cv_status::timeout) {
            std::stringstream ss;
            ss << "ThriftServer '" << _thrift_server->_name
               << "' (on port: " << _thrift_server->_port << ") did not start within " << TIMEOUT_MS
               << "ms";
            LOG(ERROR) << ss.str();
            return Status::InternalError(ss.str());
        }
    }

    // _started == true only if preServe was called. May be false if there was an exception
    // after preServe that was caught by Supervise, causing it to reset the error condition.
    if (!_thrift_server->_started) {
        std::stringstream ss;
        ss << "ThriftServer '" << _thrift_server->_name << "' (on port: " << _thrift_server->_port
           << ") did not start correctly ";
        LOG(ERROR) << ss.str();
        return Status::InternalError(ss.str());
    }

    return Status::OK();
}

void ThriftServer::ThriftServerEventProcessor::supervise() {
    DCHECK(_thrift_server->_server.get() != nullptr);

    try {
        _thrift_server->_server->serve();
    } catch (apache::thrift::TException& e) {
        LOG(ERROR) << "ThriftServer '" << _thrift_server->_name
                   << "' (on port: " << _thrift_server->_port
                   << ") exited due to TException: " << e.what();
    }

    LOG(INFO) << "ThriftServer " << _thrift_server->_name << " exited";

    {
        // _signal_lock ensures mutual exclusion of access to _thrift_server
        std::lock_guard<std::mutex> lock(_signal_lock);
        _thrift_server->_started = false;

        // There may not be anyone waiting on this signal (if the
        // exception occurs after startup). That's not a problem, this is
        // just to avoid waiting for the timeout in case of a bind
        // failure, for example.
        _signal_fired = true;
    }

    _signal_cond.notify_all();
}

void ThriftServer::ThriftServerEventProcessor::preServe() {
    // Acquire the signal lock to ensure that StartAndWaitForServer is
    // waiting on _signal_cond when we notify.
    std::lock_guard<std::mutex> lock(_signal_lock);
    _signal_fired = true;

    // This is the (only) success path - if this is not reached within TIMEOUT_MS,
    // StartAndWaitForServer will indicate failure.
    _thrift_server->_started = true;

    // Should only be one thread waiting on _signal_cond, but wake all just in case.
    _signal_cond.notify_all();
}

// This thread-local variable contains the current session key for whichever
// thrift server is currently serving a request on the current thread.
__thread ThriftServer::SessionKey* _session_key;

ThriftServer::SessionKey* ThriftServer::get_thread_session_key() {
    return _session_key;
}

void* ThriftServer::ThriftServerEventProcessor::createContext(
        std::shared_ptr<apache::thrift::protocol::TProtocol> input,
        std::shared_ptr<apache::thrift::protocol::TProtocol> output) {
    std::stringstream ss;

    apache::thrift::transport::TSocket* socket = nullptr;
    apache::thrift::transport::TTransport* transport = input->getTransport().get();
    {
        switch (_thrift_server->_server_type) {
        case NON_BLOCKING:
            socket = static_cast<apache::thrift::transport::TSocket*>(
                    static_cast<apache::thrift::transport::TFramedTransport*>(transport)
                            ->getUnderlyingTransport()
                            .get());
            break;

        case THREAD_POOL:
        case THREADED:
            socket = static_cast<apache::thrift::transport::TSocket*>(
                    static_cast<apache::thrift::transport::TBufferedTransport*>(transport)
                            ->getUnderlyingTransport()
                            .get());
            break;

        default:
            DCHECK(false) << "Unexpected thrift server type";
        }
    }

    ss << socket->getPeerAddress() << ":" << socket->getPeerPort();

    {
        std::lock_guard<std::mutex> _l(_thrift_server->_session_keys_lock);

        std::shared_ptr<SessionKey> key_ptr(new std::string(ss.str()));

        _session_key = key_ptr.get();
        _thrift_server->_session_keys[key_ptr.get()] = key_ptr;
    }

    if (_thrift_server->_session_handler != nullptr) {
        _thrift_server->_session_handler->session_start(*_session_key);
    }

    _thrift_server->thrift_connections_total->increment(1L);
    _thrift_server->thrift_current_connections->increment(1L);

    // Store the _session_key in the per-client context to avoid recomputing
    // it. If only this were accessible from RPC method calls, we wouldn't have to
    // mess around with thread locals.
    return (void*)_session_key;
}

void ThriftServer::ThriftServerEventProcessor::processContext(
        void* context, std::shared_ptr<apache::thrift::transport::TTransport> transport) {
    _session_key = reinterpret_cast<SessionKey*>(context);
}

void ThriftServer::ThriftServerEventProcessor::deleteContext(
        void* serverContext, std::shared_ptr<apache::thrift::protocol::TProtocol> input,
        std::shared_ptr<apache::thrift::protocol::TProtocol> output) {
    _session_key = (SessionKey*)serverContext;

    if (_thrift_server->_session_handler != nullptr) {
        _thrift_server->_session_handler->session_end(*_session_key);
    }

    {
        std::lock_guard<std::mutex> _l(_thrift_server->_session_keys_lock);
        _thrift_server->_session_keys.erase(_session_key);
    }

    _thrift_server->thrift_current_connections->increment(-1L);
}

ThriftServer::ThriftServer(const std::string& name,
                           const std::shared_ptr<apache::thrift::TProcessor>& processor, int port,
                           int num_worker_threads, ServerType server_type)
        : _started(false),
          _port(port),
          _num_worker_threads(num_worker_threads),
          _server_type(server_type),
          _name(name),
          _server_thread(nullptr),
          _server(nullptr),
          _processor(processor),
          _session_handler(nullptr) {
    _thrift_server_metric_entity = DorisMetrics::instance()->metric_registry()->register_entity(
            std::string("thrift_server.") + name, {{"name", name}});
    INT_GAUGE_METRIC_REGISTER(_thrift_server_metric_entity, thrift_current_connections);
    INT_COUNTER_METRIC_REGISTER(_thrift_server_metric_entity, thrift_connections_total);
}

ThriftServer::~ThriftServer() {
    stop();
    join();
}

Status ThriftServer::start() {
    DCHECK(!_started);
    std::shared_ptr<apache::thrift::protocol::TProtocolFactory> protocol_factory(
            new apache::thrift::protocol::TBinaryProtocolFactory());
    std::shared_ptr<apache::thrift::concurrency::ThreadManager> thread_mgr;
    std::shared_ptr<apache::thrift::concurrency::ThreadFactory> thread_factory =
            std::make_shared<apache::thrift::concurrency::ThreadFactory>();
    std::shared_ptr<apache::thrift::transport::TServerTransport> fe_server_transport;
    std::shared_ptr<apache::thrift::transport::TTransportFactory> transport_factory;
    std::shared_ptr<apache::thrift::transport::TNonblockingServerSocket> socket;
    if (_server_type != THREADED) {
        thread_mgr = apache::thrift::concurrency::ThreadManager::newSimpleThreadManager(
                _num_worker_threads);
        thread_mgr->threadFactory(thread_factory);
        thread_mgr->start();
    }

    // Note - if you change the transport types here, you must check that the
    // logic in createContext is still accurate.
    apache::thrift::transport::TServerSocket* server_socket = nullptr;

    switch (_server_type) {
    case NON_BLOCKING:
        socket = std::make_shared<ImprovedNonblockingServerSocket>(_port);
        if (transport_factory == nullptr) {
            transport_factory.reset(new apache::thrift::transport::TTransportFactory());
        }

        _server = std::make_unique<apache::thrift::server::TNonblockingServer>(
                _processor, transport_factory, transport_factory, protocol_factory,
                protocol_factory, socket, thread_mgr);
        break;

    case THREAD_POOL:
        fe_server_transport.reset(new apache::thrift::transport::TServerSocket(
                BackendOptions::get_service_bind_address_without_bracket(), _port));

        if (transport_factory == nullptr) {
            transport_factory = std::make_shared<ImprovedBufferedTransportFactory>();
        }

        _server = std::make_unique<apache::thrift::server::TThreadPoolServer>(
                _processor, fe_server_transport, transport_factory, protocol_factory, thread_mgr);
        break;

    case THREADED:
        server_socket = new apache::thrift::transport::TServerSocket(
                BackendOptions::get_service_bind_address_without_bracket(), _port);
        fe_server_transport.reset(server_socket);
        server_socket->setKeepAlive(true);

        if (transport_factory == nullptr) {
            transport_factory = std::make_shared<ImprovedBufferedTransportFactory>();
        }

        _server = std::make_unique<apache::thrift::server::TThreadedServer>(
                _processor, fe_server_transport, transport_factory, protocol_factory,
                thread_factory);
        break;

    default:
        std::stringstream error_msg;
        error_msg << "Unsupported server type: " << _server_type;
        LOG(ERROR) << error_msg.str();
        return Status::InternalError(error_msg.str());
    }

    std::shared_ptr<ThriftServer::ThriftServerEventProcessor> event_processor(
            new ThriftServer::ThriftServerEventProcessor(this));
    _server->setServerEventHandler(event_processor);

    RETURN_IF_ERROR(event_processor->start_and_wait_for_server());

    LOG(INFO) << "ThriftServer '" << _name << "' started on port: " << _port;

    DCHECK(_started);
    return Status::OK();
}

void ThriftServer::stop() {
    _server->stop();
}

void ThriftServer::join() {
    DCHECK(_server_thread != nullptr);
    _server_thread->join();
}

void ThriftServer::stop_for_testing() {
    DCHECK(_server_thread != nullptr);
    DCHECK(_server);
    DCHECK_EQ(_server_type, THREADED);
    _server->stop();

    if (_started) {
        join();
    }
}
} // namespace doris