/**
 *    Copyright (C) 2018-present MongoDB, Inc.
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the Server Side Public License, version 1,
 *    as published by MongoDB, Inc.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    Server Side Public License for more details.
 *
 *    You should have received a copy of the Server Side Public License
 *    along with this program. If not, see
 *    <http://www.mongodb.com/licensing/server-side-public-license>.
 *
 *    As a special exception, the copyright holders give permission to link the
 *    code of portions of this program with the OpenSSL library under certain
 *    conditions as described in each individual source file and distribute
 *    linked combinations including the program with the OpenSSL library. You
 *    must comply with the Server Side Public License in all respects for
 *    all of the code used other than as permitted herein. If you modify file(s)
 *    with this exception, you may extend this exception to your version of the
 *    file(s), but you are not obligated to do so. If you do not wish to do so,
 *    delete this exception statement from your version. If you delete this
 *    exception statement from all source files in the program, then also delete
 *    it in the license file.
 */

#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest

#include "mongo/platform/basic.h"

#include "mongo/client/dbclient_connection.h"
#include "mongo/client/dbclient_rs.h"
#include "mongo/db/query/cursor_response.h"
#include "mongo/logv2/log.h"
#include "mongo/rpc/get_status_from_command_result.h"
#include "mongo/stdx/future.h"
#include "mongo/unittest/integration_test.h"
#include "mongo/unittest/unittest.h"
#include "mongo/util/system_clock_source.h"

namespace mongo {
namespace {

class CurrentOpExhaustCursorTestFixture : public unittest::Test {
public:
    CurrentOpExhaustCursorTestFixture() = default;
    explicit CurrentOpExhaustCursorTestFixture(NetworkOp op) : _op{op} {}

    std::unique_ptr<DBClientBase> connect() {
        return connect(testAppName);
    }

    std::unique_ptr<DBClientBase> connect(StringData appName) {
        auto swConn = unittest::getFixtureConnectionString().connect(appName.toString());
        uassertStatusOK(swConn.getStatus());
        return std::move(swConn.getValue());
    }

    void initTestCollection(DBClientBase* conn) {
        // Drop and recreate the test namespace.
        conn->dropCollection(testNSS.ns());
        if (_op == dbMsg) {
            for (int i = 0; i < 10; i++) {
                auto insertCmd =
                    BSON("insert" << testNSS.coll() << "documents" << BSON_ARRAY(BSON("a" << i)));
                auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody(testNSS.db(), insertCmd));
                ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
            }
        } else if (_op == dbQuery) {
            // For dbQuery exhaust, the getMore batch size is fixed at the server-defined
            // limit of 16Mb, and can't be manipulated. We need to insert enough documents
            // to make sure we have multiple batches of getMores given that limit;
            // as a result, we create approximately 16KB documents and insert 3000 of them
            // to ensure that three getMores are required to exhaust the cursor.
            constexpr int numDocs = 3000;
            constexpr int strSize = 16 * 1024;
            for (int i = 0; i < numDocs; i++) {
                std::string bigString(strSize, 'A');
                auto insertCmd =
                    BSON("insert" << testNSS.coll() << "documents"
                                  << BSON_ARRAY(BSON("a" << i << "data" << bigString)));
                auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody(testNSS.db(), insertCmd));
                ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
            }
        }
    }

    void setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(DBClientBase* conn, bool enable) {
        auto cmdObj =
            BSON("configureFailPoint"
                 << "waitWithPinnedCursorDuringGetMoreBatch"
                 << "mode" << (enable ? "alwaysOn" : "off") << "data"
                 << BSON("shouldNotdropLock" << true << "shouldContinueOnInterrupt" << true));
        auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody("admin", cmdObj));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
    }

    std::string setWaitAfterGetMoreFinishesExecutionFailpoint(DBClientBase* conn, bool enable) {
        // Since one uses the command path and the other uses the OP_GET_MORE path, we
        // need to select the appropriate post-get-more execution fail point depending on the
        // op we are using.
        std::string fpName = (_op == dbMsg) ? "waitAfterCommandFinishesExecution"
                                            : "waitAfterOpGetMoreFinishesExecution";
        auto cmdObj = BSON("configureFailPoint" << fpName << "mode" << (enable ? "alwaysOn" : "off")
                                                << "data" << BSON("ns" << testNSS.toString()));
        auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody("admin", cmdObj));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
        // return the appropriate failPointMsg for the failpoint we chose.
        return fpName;
    }

    void setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(DBClientBase* conn,
                                                                          bool enable) {
        auto cmdObj = BSON("configureFailPoint"
                           << "waitBeforeUnpinningOrDeletingCursorAfterGetMoreBatch"
                           << "mode" << (enable ? "alwaysOn" : "off"));
        auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody("admin", cmdObj));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
    }

    bool confirmCurrentOpContents(DBClientBase* conn,
                                  BSONObj curOpMatch,
                                  bool expectEmptyResult = false,
                                  Milliseconds timeoutMS = Milliseconds(5 * 60 * 1000),
                                  Milliseconds intervalMS = Milliseconds(200)) {
        auto curOpCmd =
            BSON("aggregate" << 1 << "cursor" << BSONObj() << "pipeline"
                             << BSON_ARRAY(BSON("$currentOp" << BSON("idleCursors" << true))
                                           << BSON("$match" << curOpMatch)));
        const auto startTime = clock->now();
        while (clock->now() - startTime < timeoutMS) {
            auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody("admin", curOpCmd));
            auto swCursorRes = CursorResponse::parseFromBSON(reply->getCommandReply());
            ASSERT_OK(swCursorRes.getStatus());
            if (swCursorRes.getValue().getBatch().empty() == expectEmptyResult) {
                return true;
            }
            sleepFor(intervalMS);
        }
        auto currentOp = BSON("currentOp" << BSON("idleCursors" << true));
        LOGV2(20606,
              "confirmCurrentOpContents fails with curOpMatch: {curOpMatch} currentOp: "
              "{conn_runCommand_OpMsgRequest_fromDBAndBody_admin_currentOp_getCommandReply}",
              "curOpMatch"_attr = curOpMatch,
              "conn_runCommand_OpMsgRequest_fromDBAndBody_admin_currentOp_getCommandReply"_attr =
                  conn->runCommand(OpMsgRequest::fromDBAndBody("admin", currentOp))
                      ->getCommandReply());
        return false;
    }

    repl::OpTime getLastAppliedOpTime(DBClientBase* conn) {
        auto reply =
            conn->runCommand(OpMsgRequest::fromDBAndBody("admin", BSON("replSetGetStatus" << 1)));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
        auto lastAppliedOpTime = reply->getCommandReply()["optimes"]["appliedOpTime"];
        return repl::OpTime(lastAppliedOpTime["ts"].timestamp(),
                            lastAppliedOpTime["t"].numberLong());
    }

    // Start an exhaust request with a batchSize of 2 in the initial 'find' and a batchSize of 1 in
    // subsequent 'getMore's.
    auto startExhaustQuery(
        DBClientBase* queryConnection,
        std::unique_ptr<DBClientCursor>& queryCursor,
        int queryOptions = 0,
        Milliseconds awaitDataTimeoutMS = Milliseconds(5000),
        const boost::optional<repl::OpTime>& lastKnownCommittedOpTime = boost::none) {
        queryOptions = queryOptions | QueryOption_Exhaust;
        if (_op == dbQuery) {
            queryOptions |= DBClientCursor::QueryOptionLocal_forceOpQuery;
        }
        auto queryThread = stdx::async(
            stdx::launch::async,
            [&queryCursor,
             queryConnection,
             queryOptions,
             awaitDataTimeoutMS,
             lastKnownCommittedOpTime,
             this] {
                // Issue the initial 'find' with a batchSize of 2 and the exhaust flag set. We then
                // iterate through the first batch and confirm that the results are as expected.
                queryCursor = queryConnection->query(testNSS, {}, 0, 0, nullptr, queryOptions, 2);
                for (int i = 0; i < 2; ++i) {
                    auto doc = queryCursor->nextSafe();
                    ASSERT(!doc["a"].eoo());
                    ASSERT(doc["a"].safeNumberInt() == i);
                }
                // Having exhausted the two results returned by the initial find, we set the
                // batchSize to 1 and issue a single getMore via DBClientCursor::more(). Because the
                // 'exhaust' flag is set, the server will generate a series of internal getMores and
                // stream them back to the client until the cursor is exhausted, without the client
                // sending any further getMore requests. We expect this request to hang at the
                // 'waitWithPinnedCursorDuringGetMoreBatch' failpoint.
                queryCursor->setBatchSize(1);
                if ((queryOptions & QueryOption_CursorTailable) &&
                    (queryOptions & QueryOption_AwaitData)) {
                    queryCursor->setAwaitDataTimeoutMS(awaitDataTimeoutMS);
                    if (lastKnownCommittedOpTime) {
                        auto term = lastKnownCommittedOpTime.get().getTerm();
                        queryCursor->setCurrentTermAndLastCommittedOpTime(term,
                                                                          lastKnownCommittedOpTime);
                    }
                }
                ASSERT(queryCursor->more());
            });

        // Wait until the parallel operation initializes its cursor.
        const auto startTime = clock->now();
        while (!queryCursor && (clock->now() - startTime < parallelWaitTimeoutMS)) {
            sleepFor(Milliseconds(10));
        }
        ASSERT(queryCursor);
        LOGV2(20607,
              "Started exhaust query with cursorId: {queryCursor_getCursorId}",
              "queryCursor_getCursorId"_attr = queryCursor->getCursorId());
        return queryThread;
    }

    void runOneGetMore(DBClientBase* conn,
                       const std::unique_ptr<DBClientCursor>& queryCursor,
                       boost::optional<int> nDocsReturned) {
        BSONObjBuilder bob;
        bob.append("command.collection", testNSS.coll());
        bob.append("command.getMore", queryCursor->getCursorId());
        bob.append("failpointMsg", "waitWithPinnedCursorDuringGetMoreBatch");
        if (nDocsReturned) {
            bob.append("cursor.nDocsReturned", *nDocsReturned);
        }
        const auto curOpMatch = bob.done();
        // Confirm that the initial getMore appears in the $currentOp output.
        ASSERT(confirmCurrentOpContents(conn, curOpMatch));

        // Airlock the failpoint by releasing it only after we enable a post-getMore failpoint. This
        // ensures that no subsequent getMores can run before we re-enable the original failpoint.
        setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(conn, true);
        setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn, false);
        // Confirm that the getMore completed its batch and hit the post-getMore failpoint.
        ASSERT(confirmCurrentOpContents(
            conn,
            BSON("command.getMore" << queryCursor->getCursorId() << "failpointMsg"
                                   << "waitBeforeUnpinningOrDeletingCursorAfterGetMoreBatch")));

        // Re-enable the original failpoint to catch the next getMore, and release the current one.
        setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn, true);
        setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(conn, false);
    }

    // Test exhaust cursor is cleaned up on client disconnect. By default, the test client
    // disconnects while the exhaust getMore is running. If disconnectAfterGetMoreBatch is set to
    // true, the test client disconnects after the exhaust getMore is run but before the server
    // sends out the response.
    void testClientDisconnect(bool disconnectAfterGetMoreBatch) {
        auto conn = connect();

        // We need to set failpoints around getMore which cause it to hang, so only test against a
        // single server rather than a replica set or mongoS. Mongos doesn't support exhuast at all
        // and the failpoints might disrupt the completion of other commands (like monitoring hellos
        // and heartbeats) for replica sets.
        if (conn->isReplicaSetMember() || conn->isMongos()) {
            return;
        }

        initTestCollection(conn.get());

        // Enable a failpoint to block getMore during execution.
        setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), true);

        const auto connStr = unittest::getFixtureConnectionString();
        const auto queryConnection = std::make_unique<DBClientConnection>();
        uassertStatusOK(
            queryConnection->connect(connStr.getServers()[0], testBackgroundAppName, boost::none));
        std::unique_ptr<DBClientCursor> queryCursor;

        // Execute a query on a separate thread, with the 'exhaust' flag set.
        auto queryThread = startExhaustQuery(queryConnection.get(), queryCursor);
        // Ensure that, regardless of whether the test completes or fails, we release all
        // failpoints.
        ON_BLOCK_EXIT([&conn, &queryThread, this] {
            setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), false);
            setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(conn.get(), false);
            queryThread.wait();
        });

        // This will allow the initial getMore to run.
        if (_op == dbQuery) {
            // We can't know how many docs to expect when using dbQuery
            // because the batch size is determined by the server.
            runOneGetMore(conn.get(), queryCursor, boost::none);
        } else {
            runOneGetMore(conn.get(), queryCursor, 2);
        }

        // The next getMore will be an exhaust getMore. Confirm that the exhaust getMore appears in
        // the $currentOp output.
        BSONObjBuilder bob;
        bob.append("command.collection", testNSS.coll());
        bob.append("command.getMore", queryCursor->getCursorId());
        bob.append("failpointMsg", "waitWithPinnedCursorDuringGetMoreBatch");
        if (_op == dbMsg) {
            // We can't know how many docs to expect when using dbQuery
            // because the batch size is determined by the server.
            bob.append("cursor.nDocsReturned", 3);
        }
        auto curOpMatch = bob.done();
        ASSERT(confirmCurrentOpContents(conn.get(), curOpMatch));

        if (disconnectAfterGetMoreBatch) {
            // Allow the exhaust getMore to run but block it before sending out the response.
            std::string failPointMsg =
                setWaitAfterGetMoreFinishesExecutionFailpoint(conn.get(), true);
            setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), false);
            ASSERT(confirmCurrentOpContents(conn.get(),
                                            BSON("command.getMore" << queryCursor->getCursorId()
                                                                   << "failpointMsg"
                                                                   << failPointMsg)));
        }

        // Kill the client connection while the exhaust getMore is blocked on the failpoint.
        queryConnection->shutdownAndDisallowReconnect();
        LOGV2(20608, "Killed exhaust connection.");

        if (disconnectAfterGetMoreBatch) {
            // Disable the failpoint to allow the exhaust getMore to continue sending out the
            // response after the client disconnects. This will result in a broken pipe error.
            setWaitAfterGetMoreFinishesExecutionFailpoint(conn.get(), false);
        }

        curOpMatch = BSON("command.collection" << testNSS.coll() << "command.getMore"
                                               << queryCursor->getCursorId());
        // Confirm that the exhaust getMore was interrupted and does not appear in the $currentOp
        // output.
        const bool expectEmptyResult = true;
        ASSERT(confirmCurrentOpContents(conn.get(), curOpMatch, expectEmptyResult));

        setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), false);
        setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(conn.get(), false);

        curOpMatch = BSON("type"
                          << "idleCursor"
                          << "cursor.cursorId" << queryCursor->getCursorId());
        // Confirm that the cursor was cleaned up and does not appear in the $currentOp idleCursor
        // output.
        ASSERT(confirmCurrentOpContents(conn.get(), curOpMatch, expectEmptyResult));
    }

    // Specifies the amount of time we are willing to wait for a parallel operation to complete.
    Milliseconds parallelWaitTimeoutMS = Milliseconds(5 * 60 * 1000);

    // Obtain a pointer to the global system clock. Used to enforce timeouts in the parallel thread.
    SystemClockSource* clock = SystemClockSource::get();

    const NamespaceString testNSS{"exhaust_cursor_currentop.testColl"};

    const StringData testAppName = "curop_exhaust_cursor_test";

    const StringData testBackgroundAppName = "curop_exhaust_cursor_test_bg";

    NetworkOp _op{dbMsg};
};
}  // namespace

TEST_F(CurrentOpExhaustCursorTestFixture, CanSeeEachExhaustCursorPseudoGetMoreInCurrentOpOutput) {
    auto conn = connect();

    // We need to set failpoints around getMore which cause it to hang, so only test against a
    // single server rather than a replica set or mongoS.
    if (conn->isReplicaSetMember() || conn->isMongos()) {
        return;
    }

    initTestCollection(conn.get());

    // Enable a failpoint to block getMore during execution.
    setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), true);

    const auto queryConnection = connect(testBackgroundAppName);
    std::unique_ptr<DBClientCursor> queryCursor;

    // Execute a query on a separate thread, with the 'exhaust' flag set.
    auto queryThread = startExhaustQuery(queryConnection.get(), queryCursor);
    // Ensure that, regardless of whether the test completes or fails, we release all failpoints.
    ON_BLOCK_EXIT([&conn, &queryThread, this] {
        setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn.get(), false);
        setWaitBeforeUnpinningOrDeletingCursorAfterGetMoreBatchFailpoint(conn.get(), false);
        queryThread.wait();
    });

    // We expect that the server, having received the first {batchSize: 1} getMore for the parallel
    // thread's exhaust cursor, will produce a series of pseudo-getMores internally and stream the
    // results back to the client until the cursor is exhausted. Here, we verify that each of these
    // pseudo-getMores appear in the $currentOp output.
    for (int i = 2; i < 10; ++i) {
        runOneGetMore(conn.get(), queryCursor, i);
    }
}


TEST_F(CurrentOpExhaustCursorTestFixture, InterruptExhaustCursorPseudoGetMoreOnClientDisconnect) {
    // Test that an exhaust getMore is interrupted on client disconnect.
    testClientDisconnect(false /* disconnectAfterGetMoreBatch */);
}

TEST_F(CurrentOpExhaustCursorTestFixture, CleanupExhaustCursorOnBrokenPipe) {
    // Test that exhaust cursor is cleaned up on broken pipe even if the exhaust getMore succeeded.
    testClientDisconnect(true /* disconnectAfterGetMoreBatch */);
}

TEST_F(CurrentOpExhaustCursorTestFixture,
       InterruptExhaustCursorPseudoGetMoreOnClientDisconnectDbQuery) {
    _op = dbQuery;
    // Test that an exhaust getMore is interrupted on client disconnect.
    testClientDisconnect(false /* disconnectAfterGetMoreBatch */);
}

TEST_F(CurrentOpExhaustCursorTestFixture, CleanupExhaustCursorOnBrokenPipeDbQuery) {
    _op = dbQuery;
    // Test that exhaust cursor is cleaned up on broken pipe even if the exhaust getMore succeeded.
    testClientDisconnect(true /* disconnectAfterGetMoreBatch */);
}

TEST_F(CurrentOpExhaustCursorTestFixture, ExhaustCursorUpdatesLastKnownCommittedOpTime) {
    auto fixtureConn = connect();

    // We need to test the lastKnownCommittedOpTime in exhaust getMore requests. So we need a
    // replica set.
    if (!fixtureConn->isReplicaSetMember()) {
        return;
    }

    // Connect directly to the primary.
    DBClientBase* conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->primaryConn();
    ASSERT(conn);

    conn->dropCollection(testNSS.ns());
    // Create a capped collection to run tailable awaitData queries on.
    conn->createCollection(testNSS.ns(),
                           1024 /* size of collection */,
                           true /* capped */,
                           10 /* max number of objects */);
    // Insert initial records into the capped collection.
    for (int i = 0; i < 5; i++) {
        auto insertCmd =
            BSON("insert" << testNSS.coll() << "documents" << BSON_ARRAY(BSON("a" << i)));
        auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody(testNSS.db(), insertCmd));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
    }

    // Get the lastAppliedOpTime after the initial inserts.
    auto lastAppliedOpTime = getLastAppliedOpTime(conn);

    // Create a new connection to the primary for the exhaust query.
    const auto fixtureQueryConn = connect(testBackgroundAppName);
    DBClientBase* queryConn =
        &static_cast<DBClientReplicaSet*>(fixtureQueryConn.get())->primaryConn();
    std::unique_ptr<DBClientCursor> queryCursor;

    // Enable a failpoint to block getMore during execution to avoid races between getCursorId() and
    // receiving new batches.
    setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn, true);

    // Initiate a tailable awaitData exhaust cursor with lastKnownCommittedOpTime being the
    // lastAppliedOpTime.
    auto queryThread = startExhaustQuery(queryConn,
                                         queryCursor,
                                         QueryOption_CursorTailable | QueryOption_AwaitData,
                                         Milliseconds(1000),  // awaitData timeout
                                         lastAppliedOpTime);  // lastKnownCommittedOpTime

    // Assert non-zero cursorId.
    auto cursorId = queryCursor->getCursorId();
    ASSERT_NE(cursorId, 0LL);

    // Disable failpoint and allow exhaust queries to run.
    setWaitWithPinnedCursorDuringGetMoreBatchFailpoint(conn, false);

    ON_BLOCK_EXIT([&conn, &queryThread] { queryThread.wait(); });

    // Test that the cursor's lastKnownCommittedOpTime is eventually advanced to the
    // lastAppliedOpTime.
    auto curOpMatch =
        BSON("command.collection" << testNSS.coll() << "command.getMore" << cursorId
                                  << "cursor.lastKnownCommittedOpTime" << lastAppliedOpTime);
    ASSERT(confirmCurrentOpContents(conn, curOpMatch));

    // Inserting more records to unblock awaitData and advance the commit point.
    for (int i = 5; i < 8; i++) {
        auto insertCmd =
            BSON("insert" << testNSS.coll() << "documents" << BSON_ARRAY(BSON("a" << i)));
        auto reply = conn->runCommand(OpMsgRequest::fromDBAndBody(testNSS.db(), insertCmd));
        ASSERT_OK(getStatusFromCommandResult(reply->getCommandReply()));
    }

    // Get the new lastAppliedOpTime after the inserts.
    lastAppliedOpTime = getLastAppliedOpTime(conn);

    // Test that the cursor's lastKnownCommittedOpTime is eventually advanced to the
    // new lastAppliedOpTime.
    curOpMatch =
        BSON("command.collection" << testNSS.coll() << "command.getMore" << cursorId
                                  << "cursor.lastKnownCommittedOpTime" << lastAppliedOpTime);
    ASSERT(confirmCurrentOpContents(conn, curOpMatch));
}
}  // namespace mongo
