engine.cpp

#include <algorithm>
#include <cstdint>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <shared_mutex>
#include <thread>
#include <unordered_map>

#include "engine.hpp"
#include "io.hpp"
#include "orderbook.hpp"

void Engine::accept(ClientConnection connection) {
  auto thread =
      std::thread(&Engine::connection_thread, this, std::move(connection));
  thread.detach();
}

bool Engine::tryMatch(std::shared_ptr<Order> activeOrder, Side &opSide) {
  // (Non exclusive) Access to side
  std::shared_lock<std::shared_mutex> sideLock{opSide.sideMutex};

  // Shared logic, whether it is selling or buying
  auto sharedLogic = [&](auto it) {
    PriceLevel &priceLevel = it->second;

    // Exclusive access to price level
    std::unique_lock<std::mutex> priceLevelLock{priceLevel.priceLevelsMutex};

    // Loop through all orders in price levels
    for (auto it = priceLevel.orders.begin();
         it != priceLevel.orders.end() && activeOrder->count > 0;) {
      std::shared_ptr<Order> order = *it;

      // Different scope to make sure shared_ptr don't destruct before lock
      {
        // Exclusive access to order
        std::unique_lock<std::mutex> orderLock{order->orderMutex};

        if (order->cancelled) {
          it = priceLevel.orders.erase(it);
          continue;
        }

        uint32_t amountTraded = std::min(activeOrder->count, order->count);
        activeOrder->count -= amountTraded;
        order->count -= amountTraded;
        uint32_t execId = order->execCount++;
        intmax_t ts = getCurrentTimestamp();

        Output::OrderExecuted(order->orderId, activeOrder->orderId, execId,
                              order->price, amountTraded, ts);

        if (order->count == 0) {
          it = priceLevel.orders.erase(it);
          continue;
        } else {
          // If order count is not 0, activeOrder is
          break;
        }
      }
    }
  };

  if (activeOrder->isSell) {
    for (auto it = opSide.priceLevels.rbegin();
         it != opSide.priceLevels.rend() && it->first >= activeOrder->price &&
         activeOrder->count > 0;
         ++it)
      sharedLogic(it);
  } else {
    for (auto it = opSide.priceLevels.begin();
         it != opSide.priceLevels.end() && it->first <= activeOrder->price &&
         activeOrder->count > 0;
         ++it)
      sharedLogic(it);
  }
  // Loop through all price levels, stop when
  // 1, Hit the end
  // 2, Resting order price > active order price
  // 3, Active order run out of amount

  return activeOrder->count == 0;
}

void Engine::addOrder(std::shared_ptr<Order> activeOrder, Side &curSide) {
  // Exclusive access to side
  std::unique_lock<std::shared_mutex> sideLock{curSide.sideMutex};

  auto it = curSide.priceLevels.find(activeOrder->price);
  if (it == curSide.priceLevels.end()) {
    PriceLevel &newPriceLevel = curSide.priceLevels[activeOrder->price];
    newPriceLevel.orders.push_back(activeOrder);
  } else {
    std::lock_guard<std::mutex> guard{it->second.priceLevelsMutex};
    it->second.orders.push_back(activeOrder);
  }
};

InstrumentOrderBook &Engine::getInstrumentBook(const std::string &instrument) {
  std::shared_lock<std::shared_mutex> lock{instrumentsMutex};
  auto it = instruments.find(instrument);
  if (it != instruments.end()) {
    return it->second;
  }
  lock.unlock();

  std::unique_lock<std::shared_mutex> ulock{instrumentsMutex};
  return instruments[instrument];
}

bool Engine::processOrder(ClientCommand cmd, intmax_t ts) {
  std::string instrument{cmd.instrument};

  auto &book = getInstrumentBook(instrument);

  bool isSell = cmd.type == input_sell;
  auto &opSide = isSell ? book.buyOrders : book.sellOrders;
  auto &sameSide = isSell ? book.sellOrders : book.buyOrders;

  std::shared_ptr<Order> activeOrder =
      std::make_shared<Order>(cmd.order_id, cmd.price, cmd.count, isSell, ts);

  bool isAdded = false;

  // Fast path: try to match against resting orders while only holding a
  // shared (read) lock on the book, so many threads can match concurrently.
  {
    std::shared_lock<std::shared_mutex> bookLock{book.bookMutex};
    if (tryMatch(activeOrder, opSide)) {
      return isAdded; // fully filled, nothing to add
    }
  }

  // Slow path: the order wasn't fully filled, so it needs to rest on the
  // book. Acquire the unique (write) lock *blocking* rather than spinning
  // with try_lock — this avoids burning CPU under contention. Because the
  // book may have changed between releasing the shared lock and acquiring
  // the unique lock (another thread could have added liquidity), re-attempt
  // the match once more while holding the write lock before inserting.
  std::unique_lock<std::shared_mutex> writeBookLock{book.bookMutex};

  if (tryMatch(activeOrder, opSide)) {
    return isAdded; // got fully filled in the window before we got the lock
  }

  {
    std::unique_lock<std::mutex> lock{ordersMutex};
    orders[activeOrder->orderId] = activeOrder;
  }
  addOrder(activeOrder, sameSide);
  isAdded = true;
  Output::OrderAdded(activeOrder->orderId, cmd.instrument, activeOrder->price,
                     activeOrder->count, isSell, getCurrentTimestamp());

  return isAdded;
};

void Engine::connection_thread(ClientConnection connection) {
  // Don't need mutex for set cuz it's only belong to 1 thread
  std::set<uint32_t> localOrders;
  while (true) {
    ClientCommand cmd{};
    switch (connection.readInput(cmd)) {
    case ReadResult::Error:
      SyncCerr{} << "Error reading input" << std::endl;
    case ReadResult::EndOfFile:
      return;
    case ReadResult::Success:
      break;
    }

    switch (cmd.type) {
    case input_cancel: {
      std::unique_lock<std::mutex> globalOrderLock{ordersMutex};
      auto it = orders.find(cmd.order_id);

      if (it == orders.end()) {
        Output::OrderDeleted(cmd.order_id, false, getCurrentTimestamp());
        break;
      }

      if (localOrders.find(cmd.order_id) == localOrders.end()) {
        Output::OrderDeleted(cmd.order_id, false, getCurrentTimestamp());
        break;
      }

      std::shared_ptr<Order> order = it->second;
      std::unique_lock<std::mutex> orderLock{order->orderMutex};
      globalOrderLock.unlock();

      if (order->cancelled || order->count == 0) {
        Output::OrderDeleted(cmd.order_id, false, getCurrentTimestamp());
        break;
      }

      order->cancelled = true;
      Output::OrderDeleted(cmd.order_id, true, getCurrentTimestamp());
      break;
    }

    case input_buy:
    case input_sell: {
      bool isAdded = processOrder(cmd, getCurrentTimestamp());
      if (isAdded) {
        localOrders.insert(cmd.order_id);
      }
      break;
    }

    default:
      break;
    }
  }
}