sockutils.hpp

#pragma once
 
#include <assert.h>
#include <span>
#include <algorithm>
#include "corochain.hpp"
 
#ifdef min
#undef min
#endif
 
#ifdef max
#undef max
#endif
 
namespace NNet {
 
struct TLine {
    std::string_view Part1;
    std::string_view Part2;
 
    size_t Size() const {
        return Part1.size() + Part2.size();
    }
 
    operator bool() const {
        return !Part1.empty();
    }
};
 
template<typename TSocket>
struct TByteReader {
    TByteReader(TSocket& socket)
        : Socket(socket)
    { }
    TFuture<void> Read(void* data, size_t size) {
        char* p = static_cast<char*>(data);
 
        if (!Buffer.empty()) {
            size_t toCopy = std::min(size, Buffer.size());
            std::copy(Buffer.begin(), Buffer.begin() + toCopy, p);
            p += toCopy;
            size -= toCopy;
            Buffer.erase(Buffer.begin(), Buffer.begin() + toCopy);
        }
 
        while (size != 0) {
            auto readSize = co_await Socket.ReadSome(p, size);
            if (readSize == 0) {
                throw std::runtime_error("Connection closed");
            }
            if (readSize < 0) {
                continue; // retry
            }
            p += readSize;
            size -= readSize;
        }
        co_return;
    }
 
    TFuture<ssize_t> ReadSome(void* data, size_t size) {
        char tempBuffer[1024];
        // read by chunks, copy to user up-to size bytes
 
        if (!Buffer.empty()) {
            size_t toCopy = std::min(size, Buffer.size());
            std::copy(Buffer.begin(), Buffer.begin() + toCopy, static_cast<char*>(data));
            Buffer.erase(Buffer.begin(), Buffer.begin() + toCopy);
            co_return toCopy;
        }
 
        auto readSize = co_await Socket.ReadSome(tempBuffer, sizeof(tempBuffer));
        if (readSize == 0) {
            co_return 0;
        }
        if (readSize < 0) {
            co_return -1; // retry
        }
 
        size_t toCopy = std::min(static_cast<size_t>(readSize), size);
        std::copy(tempBuffer, tempBuffer + toCopy, static_cast<char*>(data));
        if (static_cast<size_t>(readSize) > toCopy) {
            Buffer.insert(Buffer.end(), tempBuffer + toCopy, tempBuffer + readSize);
        }
        co_return toCopy;
    }
 
    TFuture<std::string> ReadUntil(const std::string& delimiter)
    {
        std::string result;
        char tempBuffer[1024];
 
        while (true) {
            auto pos = std::search(Buffer.begin(), Buffer.end(), delimiter.begin(), delimiter.end());
            if (pos != Buffer.end()) {
                size_t delimiterOffset = std::distance(Buffer.begin(), pos);
                result.insert(result.end(), Buffer.begin(), Buffer.begin() + delimiterOffset + delimiter.size());
                Buffer.erase(Buffer.begin(), Buffer.begin() + delimiterOffset + delimiter.size());
                co_return result;
            }
 
            result.insert(result.end(), Buffer.begin(), Buffer.end());
            Buffer.clear();
 
            auto readSize = co_await Socket.ReadSome(tempBuffer, sizeof(tempBuffer));
            if (readSize == 0) {
                throw std::runtime_error("Connection closed");
            }
            if (readSize < 0) {
                continue; // retry
            }
 
            Buffer.insert(Buffer.end(), tempBuffer, tempBuffer + readSize);
        }
 
        co_return result;
    }
 
private:
    TSocket& Socket;
    std::deque<char> Buffer;
};
 
template<typename TSocket>
struct TByteWriter {
    TByteWriter(TSocket& socket)
        : Socket(socket)
    { }
    TFuture<void> Write(const void* data, size_t size) {
        const char* p = static_cast<const char*>(data);
        while (size != 0) {
            auto readSize = co_await Socket.WriteSome(p, size);
            if (readSize == 0) {
                throw std::runtime_error("Connection closed");
            }
            if (readSize < 0) {
                continue; // retry
            }
            p += readSize;
            size -= readSize;
        }
        co_return;
    }
    TFuture<void> Write(const TLine& line) {
        co_await Write(line.Part1.data(), line.Part1.size());
        co_await Write(line.Part2.data(), line.Part2.size());
        co_return;
    }
 
private:
    TSocket& Socket;
};
 
template<typename T, typename TSocket>
struct TStructReader {
    TStructReader(TSocket& socket)
        : Socket(socket)
    { }
    TFuture<T> Read() {
        T res;
        size_t size = sizeof(T);
        char* p = reinterpret_cast<char*>(&res);
        while (size != 0) {
            auto readSize = co_await Socket.ReadSome(p, size);
            if (readSize == 0) {
                throw std::runtime_error("Connection closed");
            }
            if (readSize < 0) {
                continue; // retry
            }
            p += readSize;
            size -= readSize;
        }
        co_return res;
    }
 
private:
    TSocket& Socket;
};
 
struct TLineSplitter {
public:
    TLineSplitter(int maxLen);
    TLine Pop();
    void Push(const char* buf, size_t size);
 
private:
    size_t WPos;
    size_t RPos;
    size_t Size;
    size_t Cap;
    std::string Data;
    std::string_view View;
};
 
struct TZeroCopyLineSplitter {
public:
    TZeroCopyLineSplitter(int maxLen);
    TLine Pop();
    std::span<char> Acquire(size_t size);
    void Commit(size_t size);
    void Push(const char* p, size_t len);
 
private:
    size_t WPos;
    size_t RPos;
    size_t Size;
    size_t Cap;
    std::string Data;
    std::string_view View;
};
template<typename TSocket>
struct TLineReader {
    TLineReader(TSocket& socket, int maxLineSize = 4096)
        : Socket(socket)
        , Splitter(maxLineSize)
        , ChunkSize(maxLineSize / 2)
    { }
    TFuture<TLine> Read() {
        auto line = Splitter.Pop();
        while (!line) {
            auto buf = Splitter.Acquire(ChunkSize);
            auto size = co_await Socket.ReadSome(buf.data(), buf.size());
            if (size < 0) {
                continue;
            }
            if (size == 0) {
                break;
            }
            Splitter.Commit(size);
            line = Splitter.Pop();
        }
        co_return line;
    }
 
private:
    TSocket& Socket;
    TZeroCopyLineSplitter Splitter;
    int ChunkSize;
};
 
} // namespace NNet {