orderbook_checkpoint.cpp

#include <iostream>
#include <fstream>
#include <map>
#include <string>
#include <vector>
#include <cstring>
#include <algorithm>
#include <sstream>
#include <unordered_map>
#include <chrono>
#include <filesystem>

// Price is stored as integer (fixed 4 decimal places)
using Price = int32_t;
using Volume = uint64_t;
using OrderId = uint64_t;
using SequenceNo = uint32_t;

enum class Side : char {
    BID = 'B',
    ASK = 'S'
};

enum class MessageType : char {
    ADD = 'A',
    UPDATE = 'U',
    DELETE = 'D',
    EXECUTE = 'E'
};

struct Order {
    Price price;
    Volume volume;
    Side side;
};

struct PriceLevel {
    Price price;
    Volume totalVolume;
    
    PriceLevel(Price p = 0, Volume v = 0) : price(p), totalVolume(v) {}
};

bool operator==(const PriceLevel& a, const PriceLevel& b) {
    return a.price == b.price && a.totalVolume == b.totalVolume;
}

class OrderBook {
private:
    std::unordered_map<OrderId, Order> orders;
    std::map<Price, Volume, std::greater<Price>> bids;
    std::map<Price, Volume> asks;
    
public:
    void addOrder(OrderId orderId, Price price, Volume volume, Side side) {
        orders[orderId] = {price, volume, side};
        
        if (side == Side::BID) {
            bids[price] += volume;
        } else {
            asks[price] += volume;
        }
    }
    
    void deleteOrder(OrderId orderId, Side side) {
        auto it = orders.find(orderId);
        if (it == orders.end()) return;
        
        Order& order = it->second;
        
        if (order.side != side) {
            std::cerr << "Warning: Side mismatch for order " << orderId << std::endl;
            return;
        }
        
        if (side == Side::BID) {
            bids[order.price] -= order.volume;
            if (bids[order.price] == 0) {
                bids.erase(order.price);
            }
        } else {
            asks[order.price] -= order.volume;
            if (asks[order.price] == 0) {
                asks.erase(order.price);
            }
        }
        
        orders.erase(it);
    }

    void executeOrder(OrderId orderId, Volume tradedQty, Side side) {
        auto it = orders.find(orderId);
        if (it == orders.end()) return;
        
        Order& order = it->second;
        
        if (order.side != side) {
            std::cerr << "Warning: Side mismatch for order " << orderId << std::endl;
            return;
        }
        
        if (side == Side::BID) {
            bids[order.price] -= tradedQty;
            if (bids[order.price] == 0) {
                bids.erase(order.price);
            }
        } else {
            asks[order.price] -= tradedQty;
            if (asks[order.price] == 0) {
                asks.erase(order.price);
            }
        }
        
        order.volume -= tradedQty;
        if (order.volume == 0) {
            orders.erase(it);
        }
    }

    void updateOrder(OrderId orderId, Price newPrice, Volume newVolume, Side side) {
        auto it = orders.find(orderId);
        if (it == orders.end()) {
            addOrder(orderId, newPrice, newVolume, side);
            return;
        }
        
        Order& order = it->second;
        
        if (order.side != side) {
            std::cerr << "Warning: Side mismatch for order " << orderId << std::endl;
            return;
        }
        
        Price oldPrice = order.price;
        Volume oldVolume = order.volume;
        
        if (side == Side::BID) {
            bids[oldPrice] -= oldVolume;
            if (bids[oldPrice] == 0) {
                bids.erase(oldPrice);
            }
            bids[newPrice] += newVolume;
        } else {
            asks[oldPrice] -= oldVolume;
            if (asks[oldPrice] == 0) {
                asks.erase(oldPrice);
            }
            asks[newPrice] += newVolume;
        }
        
        order.price = newPrice;
        order.volume = newVolume;
    }
    
    std::vector<PriceLevel> getTopBids(int levels) const {
        std::vector<PriceLevel> result;
        int count = 0;
        for (const auto& [price, volume] : bids) {
            if (count >= levels) break;
            result.push_back({price, volume});
            count++;
        }
        return result;
    }
    
    std::vector<PriceLevel> getTopAsks(int levels) const {
        std::vector<PriceLevel> result;
        int count = 0;
        for (const auto& [price, volume] : asks) {
            if (count >= levels) break;
            result.push_back({price, volume});
            count++;
        }
        return result;
    }
    
    // Checkpoint methods
    void saveToFile(std::ofstream& out) const {
        // Save number of orders
        size_t orderCount = orders.size();
        out.write(reinterpret_cast<const char*>(&orderCount), sizeof(orderCount));
        
        // Save each order
        for (const auto& [orderId, order] : orders) {
            out.write(reinterpret_cast<const char*>(&orderId), sizeof(orderId));
            out.write(reinterpret_cast<const char*>(&order.price), sizeof(order.price));
            out.write(reinterpret_cast<const char*>(&order.volume), sizeof(order.volume));
            out.write(reinterpret_cast<const char*>(&order.side), sizeof(order.side));
        }
    }
    
    void loadFromFile(std::ifstream& in) {
        orders.clear();
        bids.clear();
        asks.clear();
        
        // Read number of orders
        size_t orderCount;
        in.read(reinterpret_cast<char*>(&orderCount), sizeof(orderCount));
        
        // Read each order and rebuild the book
        for (size_t i = 0; i < orderCount; i++) {
            OrderId orderId;
            Order order;
            
            in.read(reinterpret_cast<char*>(&orderId), sizeof(orderId));
            in.read(reinterpret_cast<char*>(&order.price), sizeof(order.price));
            in.read(reinterpret_cast<char*>(&order.volume), sizeof(order.volume));
            in.read(reinterpret_cast<char*>(&order.side), sizeof(order.side));
            
            addOrder(orderId, order.price, order.volume, order.side);
        }
    }
};

struct Snapshot {
    std::vector<PriceLevel> bids;
    std::vector<PriceLevel> asks;
    
    bool operator==(const Snapshot& other) const {
        return bids == other.bids && asks == other.asks;
    }
    
    bool operator!=(const Snapshot& other) const {
        return !(*this == other);
    }
};

class CheckpointManager {
private:
    std::string checkpointDir;
    std::string messageLogPath;
    std::ofstream messageLog;
    SequenceNo lastCheckpointSeq;
    int checkpointInterval;
    
public:
    CheckpointManager(const std::string& dir, int interval) 
        : checkpointDir(dir), lastCheckpointSeq(0), checkpointInterval(interval) {
        
        // Create checkpoint directory if it doesn't exist
        std::filesystem::create_directories(checkpointDir);
        
        messageLogPath = checkpointDir + "/messages.log";
        messageLog.open(messageLogPath, std::ios::binary | std::ios::app);
        
        if (!messageLog.is_open()) {
            throw std::runtime_error("Failed to open message log file");
        }
    }
    
    ~CheckpointManager() {
        if (messageLog.is_open()) {
            messageLog.close();
        }
    }
    
    void logMessage(SequenceNo seqNo, const std::vector<char>& msgData) {
        uint32_t msgSize = msgData.size();
        messageLog.write(reinterpret_cast<const char*>(&seqNo), sizeof(seqNo));
        messageLog.write(reinterpret_cast<const char*>(&msgSize), sizeof(msgSize));
        messageLog.write(msgData.data(), msgSize);
        messageLog.flush();
    }
    
    bool shouldCheckpoint(SequenceNo seqNo) const {
        return (seqNo - lastCheckpointSeq) >= checkpointInterval;
    }
    
    void saveCheckpoint(SequenceNo seqNo, 
                       const std::map<std::string, OrderBook>& books,
                       const std::map<std::string, Snapshot>& snapshots) {
        
        std::string checkpointPath = checkpointDir + "/checkpoint_" + std::to_string(seqNo) + ".dat";
        std::ofstream out(checkpointPath, std::ios::binary);
        
        if (!out.is_open()) {
            std::cerr << "Failed to create checkpoint file: " << checkpointPath << std::endl;
            return;
        }
        
        // Write sequence number
        out.write(reinterpret_cast<const char*>(&seqNo), sizeof(seqNo));
        
        // Write number of symbols
        size_t symbolCount = books.size();
        out.write(reinterpret_cast<const char*>(&symbolCount), sizeof(symbolCount));
        
        // Write each symbol's data
        for (const auto& [symbol, book] : books) {
            // Write symbol length and symbol
            size_t symbolLen = symbol.length();
            out.write(reinterpret_cast<const char*>(&symbolLen), sizeof(symbolLen));
            out.write(symbol.c_str(), symbolLen);
            
            // Write order book
            book.saveToFile(out);
            
            // Write last snapshot
            auto snapIt = snapshots.find(symbol);
            if (snapIt != snapshots.end()) {
                const Snapshot& snap = snapIt->second;
                
                // Write bids
                size_t bidCount = snap.bids.size();
                out.write(reinterpret_cast<const char*>(&bidCount), sizeof(bidCount));
                for (const auto& bid : snap.bids) {
                    out.write(reinterpret_cast<const char*>(&bid.price), sizeof(bid.price));
                    out.write(reinterpret_cast<const char*>(&bid.totalVolume), sizeof(bid.totalVolume));
                }
                
                // Write asks
                size_t askCount = snap.asks.size();
                out.write(reinterpret_cast<const char*>(&askCount), sizeof(askCount));
                for (const auto& ask : snap.asks) {
                    out.write(reinterpret_cast<const char*>(&ask.price), sizeof(ask.price));
                    out.write(reinterpret_cast<const char*>(&ask.totalVolume), sizeof(ask.totalVolume));
                }
            }
        }
        
        out.close();
        lastCheckpointSeq = seqNo;
        
        // Clean up old checkpoints (keep last 3)
        cleanupOldCheckpoints();
        
        std::cout << "# Checkpoint saved at sequence " << seqNo << std::endl;
    }
    
    bool loadLatestCheckpoint(SequenceNo& seqNo,
                              std::map<std::string, OrderBook>& books,
                              std::map<std::string, Snapshot>& snapshots) {
        
        // Find latest checkpoint file
        std::vector<std::string> checkpointFiles;
        for (const auto& entry : std::filesystem::directory_iterator(checkpointDir)) {
            if (entry.path().filename().string().find("checkpoint_") == 0) {
                checkpointFiles.push_back(entry.path().string());
            }
        }
        
        if (checkpointFiles.empty()) {
            return false;
        }
        
        // Sort to get latest
        std::sort(checkpointFiles.begin(), checkpointFiles.end());
        std::string latestCheckpoint = checkpointFiles.back();
        
        std::ifstream in(latestCheckpoint, std::ios::binary);
        if (!in.is_open()) {
            std::cerr << "Failed to open checkpoint: " << latestCheckpoint << std::endl;
            return false;
        }
        
        // Read sequence number
        in.read(reinterpret_cast<char*>(&seqNo), sizeof(seqNo));
        
        // Read number of symbols
        size_t symbolCount;
        in.read(reinterpret_cast<char*>(&symbolCount), sizeof(symbolCount));
        
        // Read each symbol's data
        for (size_t i = 0; i < symbolCount; i++) {
            // Read symbol
            size_t symbolLen;
            in.read(reinterpret_cast<char*>(&symbolLen), sizeof(symbolLen));
            
            std::vector<char> symbolBuf(symbolLen);
            in.read(symbolBuf.data(), symbolLen);
            std::string symbol(symbolBuf.begin(), symbolBuf.end());
            
            // Read order book
            books[symbol].loadFromFile(in);
            
            // Read snapshot
            Snapshot snap;
            
            // Read bids
            size_t bidCount;
            in.read(reinterpret_cast<char*>(&bidCount), sizeof(bidCount));
            for (size_t j = 0; j < bidCount; j++) {
                Price price;
                Volume volume;
                in.read(reinterpret_cast<char*>(&price), sizeof(price));
                in.read(reinterpret_cast<char*>(&volume), sizeof(volume));
                snap.bids.push_back({price, volume});
            }
            
            // Read asks
            size_t askCount;
            in.read(reinterpret_cast<char*>(&askCount), sizeof(askCount));
            for (size_t j = 0; j < askCount; j++) {
                Price price;
                Volume volume;
                in.read(reinterpret_cast<char*>(&price), sizeof(price));
                in.read(reinterpret_cast<char*>(&volume), sizeof(volume));
                snap.asks.push_back({price, volume});
            }
            
            snapshots[symbol] = snap;
        }
        
        in.close();
        lastCheckpointSeq = seqNo;
        
        std::cout << "# Loaded checkpoint from sequence " << seqNo << std::endl;
        return true;
    }
    
    std::vector<std::pair<SequenceNo, std::vector<char>>> replayMessages(SequenceNo fromSeq) {
        std::vector<std::pair<SequenceNo, std::vector<char>>> messages;
        
        std::ifstream log(messageLogPath, std::ios::binary);
        if (!log.is_open()) {
            return messages;
        }
        
        while (log.good()) {
            SequenceNo seqNo;
            uint32_t msgSize;
            
            log.read(reinterpret_cast<char*>(&seqNo), sizeof(seqNo));
            if (log.gcount() != sizeof(seqNo)) break;
            
            log.read(reinterpret_cast<char*>(&msgSize), sizeof(msgSize));
            if (log.gcount() != sizeof(msgSize)) break;
            
            std::vector<char> msgData(msgSize);
            log.read(msgData.data(), msgSize);
            if (log.gcount() != static_cast<std::streamsize>(msgSize)) break;
            
            if (seqNo > fromSeq) {
                messages.push_back({seqNo, msgData});
            }
        }
        
        log.close();
        
        std::cout << "# Replaying " << messages.size() << " messages from sequence " 
                  << fromSeq + 1 << std::endl;
        
        return messages;
    }
    
private:
    void cleanupOldCheckpoints() {
        std::vector<std::string> checkpointFiles;
        for (const auto& entry : std::filesystem::directory_iterator(checkpointDir)) {
            if (entry.path().filename().string().find("checkpoint_") == 0) {
                checkpointFiles.push_back(entry.path().string());
            }
        }
        
        if (checkpointFiles.size() <= 3) {
            return;
        }
        
        std::sort(checkpointFiles.begin(), checkpointFiles.end());
        
        // Remove all but last 3
        for (size_t i = 0; i < checkpointFiles.size() - 3; i++) {
            std::filesystem::remove(checkpointFiles[i]);
        }
    }
};

class OrderBookProcessor {
private:
    std::map<std::string, OrderBook> books;
    std::map<std::string, Snapshot> lastSnapshots;
    int depthLevels;
    CheckpointManager* checkpointMgr;
    
public:
    OrderBookProcessor(int levels) : depthLevels(levels), checkpointMgr(nullptr) {}
    
    void setCheckpointManager(CheckpointManager* mgr) {
        checkpointMgr = mgr;
    }
    
    void processMessage(SequenceNo seqNo, const std::vector<char>& msgData) {
        if (msgData.empty()) return;
        
        // Log message for recovery
        if (checkpointMgr) {
            checkpointMgr->logMessage(seqNo, msgData);
        }
        
        MessageType msgType = static_cast<MessageType>(msgData[0]);
        
        switch (msgType) {
            case MessageType::ADD:
                processAdd(seqNo, msgData);
                break;
            case MessageType::UPDATE:
                processUpdate(seqNo, msgData);
                break;
            case MessageType::DELETE:
                processDelete(seqNo, msgData);
                break;
            case MessageType::EXECUTE:
                processExecute(seqNo, msgData);
                break;
            default:
                std::cerr << "Unknown message type: " << static_cast<char>(msgType) << std::endl;
                break;
        }
        
        // Save checkpoint if needed
        if (checkpointMgr && checkpointMgr->shouldCheckpoint(seqNo)) {
            checkpointMgr->saveCheckpoint(seqNo, books, lastSnapshots);
        }
    }
    
    bool restoreFromCheckpoint(CheckpointManager& mgr) {
        SequenceNo checkpointSeq;
        if (mgr.loadLatestCheckpoint(checkpointSeq, books, lastSnapshots)) {
            // Replay messages after checkpoint
            auto messages = mgr.replayMessages(checkpointSeq);
            
            for (const auto& [seqNo, msgData] : messages) {
                processMessage(seqNo, msgData);
            }
            
            return true;
        }
        return false;
    }
    
private:
    void processAdd(SequenceNo seqNo, const std::vector<char>& msg) {
        std::string symbol = readAlpha(msg, 1, 3);
        OrderId orderId = readNumeric<OrderId>(msg, 4);
        Side side = static_cast<Side>(msg[12]);
        Volume size = readNumeric<Volume>(msg, 16);
        Price price = readPrice(msg, 24);
        
        books[symbol].addOrder(orderId, price, size, side);
        checkAndPrintSnapshot(seqNo, symbol);
    }
    
    void processUpdate(SequenceNo seqNo, const std::vector<char>& msg) {
        std::string symbol = readAlpha(msg, 1, 3);
        OrderId orderId = readNumeric<OrderId>(msg, 4);
        Side side = static_cast<Side>(msg[12]);
        Volume size = readNumeric<Volume>(msg, 16);
        Price price = readPrice(msg, 24);
        
        books[symbol].updateOrder(orderId, price, size, side);
        checkAndPrintSnapshot(seqNo, symbol);
    }
    
    void processDelete(SequenceNo seqNo, const std::vector<char>& msg) {
        std::string symbol = readAlpha(msg, 1, 3);
        OrderId orderId = readNumeric<OrderId>(msg, 4);
        Side side = static_cast<Side>(msg[12]);
        
        books[symbol].deleteOrder(orderId, side);
        checkAndPrintSnapshot(seqNo, symbol);
    }
    
    void processExecute(SequenceNo seqNo, const std::vector<char>& msg) {
        std::string symbol = readAlpha(msg, 1, 3);
        OrderId orderId = readNumeric<OrderId>(msg, 4);
        Side side = static_cast<Side>(msg[12]);
        Volume tradedQty = readNumeric<Volume>(msg, 16);
        
        books[symbol].executeOrder(orderId, tradedQty, side);
        checkAndPrintSnapshot(seqNo, symbol);
    }
    
    void checkAndPrintSnapshot(SequenceNo seqNo, const std::string& symbol) {
        Snapshot current;
        current.bids = books[symbol].getTopBids(depthLevels);
        current.asks = books[symbol].getTopAsks(depthLevels);
        
        if (lastSnapshots[symbol] != current) {
            printSnapshot(seqNo, symbol, current);
            lastSnapshots[symbol] = current;
        }
    }
    
    void printSnapshot(SequenceNo seqNo, const std::string& symbol, const Snapshot& snap) {
        std::cout << seqNo << ", " << symbol << ", [";
        
        for (size_t i = 0; i < snap.bids.size(); i++) {
            if (i > 0) std::cout << ", ";
            std::cout << "(" << snap.bids[i].price << ", " << snap.bids[i].totalVolume << ")";
        }
        
        std::cout << "], [";
        
        for (size_t i = 0; i < snap.asks.size(); i++) {
            if (i > 0) std::cout << ", ";
            std::cout << "(" << snap.asks[i].price << ", " << snap.asks[i].totalVolume << ")";
        }
        
        std::cout << "]" << std::endl;
    }
    
    std::string readAlpha(const std::vector<char>& data, size_t offset, size_t len) {
        std::string result(data.begin() + offset, data.begin() + offset + len);
        result.erase(result.find_last_not_of(' ') + 1);
        return result;
    }
    
    template<typename T>
    T readNumeric(const std::vector<char>& data, size_t offset) {
        T value;
        std::memcpy(&value, &data[offset], sizeof(T));
        return value;
    }
    
    Price readPrice(const std::vector<char>& data, size_t offset) {
        Price value;
        std::memcpy(&value, &data[offset], sizeof(Price));
        return value;
    }
};

int main(int argc, char* argv[]) {
    if (argc < 2 || argc > 4) {
        std::cerr << "Usage: " << argv[0] << " <depth_levels> [checkpoint_dir] [checkpoint_interval]" << std::endl;
        std::cerr << "  checkpoint_dir: directory for checkpoints (default: ./checkpoints)" << std::endl;
        std::cerr << "  checkpoint_interval: messages between checkpoints (default: 1000)" << std::endl;
        return 1;
    }
    
    int depthLevels = std::atoi(argv[1]);
    if (depthLevels <= 0) {
        std::cerr << "Depth levels must be positive" << std::endl;
        return 1;
    }
    
    std::string checkpointDir = (argc >= 3) ? argv[2] : "./checkpoints";
    int checkpointInterval = (argc >= 4) ? std::atoi(argv[3]) : 1000;
    
    try {
        CheckpointManager checkpointMgr(checkpointDir, checkpointInterval);
        OrderBookProcessor processor(depthLevels);
        processor.setCheckpointManager(&checkpointMgr);
        
        // Try to restore from checkpoint
        if (processor.restoreFromCheckpoint(checkpointMgr)) {
            std::cout << "# Recovery completed successfully" << std::endl;
        }
        
        // Continue processing new messages from stdin
        std::cin.sync_with_stdio(false);
        
        while (std::cin.good()) {
            SequenceNo seqNo;
            uint32_t msgSize;
            
            std::cin.read(reinterpret_cast<char*>(&seqNo), sizeof(seqNo));
            if (std::cin.gcount() != sizeof(seqNo)) break;
            
            std::cin.read(reinterpret_cast<char*>(&msgSize), sizeof(msgSize));
            if (std::cin.gcount() != sizeof(msgSize)) break;
            
            std::vector<char> msgData(msgSize);
            std::cin.read(msgData.data(), msgSize);
            if (std::cin.gcount() != static_cast<std::streamsize>(msgSize)) break;
            
            processor.processMessage(seqNo, msgData);
        }
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}