aboutsummaryrefslogtreecommitdiffstats
path: root/vbench/src/vbench/core/socket.h
blob: e961994c37ae95191dcdcccf6c40b16a1a60b1c5 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#pragma once

#include "string.h"
#include "stream.h"
#include <vespa/vespalib/data/simple_buffer.h>
#include <vespa/vespalib/net/socket_handle.h>
#include <vespa/vespalib/net/server_socket.h>
#include <memory>

namespace vbench {

using Input = vespalib::Input;
using Memory = vespalib::Memory;
using Output = vespalib::Output;
using SimpleBuffer = vespalib::SimpleBuffer;
using WritableMemory = vespalib::WritableMemory;

class Socket : public Stream
{
private:
    vespalib::SocketHandle _socket;
    SimpleBuffer           _input;
    SimpleBuffer           _output;
    Taint                  _taint;
    bool                   _eof;

public:
    Socket(vespalib::SocketHandle socket);
    Socket(const string &host, int port);
    virtual bool eof() const override { return _eof; }
    virtual Memory obtain() override;
    virtual Input &evict(size_t bytes) override;
    virtual WritableMemory reserve(size_t bytes) override;
    virtual Output &commit(size_t bytes) override;
    virtual const Taint &tainted() const override { return _taint; }
};

struct ServerSocket {
    vespalib::ServerSocket server_socket;
    ServerSocket() : server_socket(0) {}
    int port() const { return server_socket.address().port(); }
    Stream::UP accept() {
        vespalib::SocketHandle handle = server_socket.accept();
        if (handle.valid()) {
            return std::make_unique<Socket>(std::move(handle));
        } else {
            return Stream::UP();
        }
    }
    void close() { server_socket.shutdown(); }
};

} // namespace vbench