108 lines
2.4 KiB
C++
108 lines
2.4 KiB
C++
|
#include "muduo/base/Logging.h"
|
||
|
|
#include "muduo/net/EventLoop.h"
|
||
|
|
#include "muduo/net/TcpServer.h"
|
||
|
|
|
||
|
|
#include "examples/wordcount/hash.h"
|
||
|
|
|
||
|
|
#include <fstream>
|
||
|
|
|
||
|
|
#include <stdio.h>
|
||
|
|
|
||
|
|
using namespace muduo;
|
||
|
|
using namespace muduo::net;
|
||
|
|
|
||
|
|
class WordCountReceiver : muduo::noncopyable
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
WordCountReceiver(EventLoop* loop, const InetAddress& listenAddr)
|
||
|
|
: loop_(loop),
|
||
|
|
server_(loop, listenAddr, "WordCountReceiver"),
|
||
|
|
senders_(0)
|
||
|
|
{
|
||
|
|
server_.setConnectionCallback(
|
||
|
|
std::bind(&WordCountReceiver::onConnection, this, _1));
|
||
|
|
server_.setMessageCallback(
|
||
|
|
std::bind(&WordCountReceiver::onMessage, this, _1, _2, _3));
|
||
|
|
}
|
||
|
|
|
||
|
|
void start(int senders)
|
||
|
|
{
|
||
|
|
LOG_INFO << "start " << senders << " senders";
|
||
|
|
senders_ = senders;
|
||
|
|
wordcounts_.clear();
|
||
|
|
server_.start();
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
void onConnection(const TcpConnectionPtr& conn)
|
||
|
|
{
|
||
|
|
LOG_DEBUG << conn->peerAddress().toIpPort() << " -> "
|
||
|
|
<< conn->localAddress().toIpPort() << " is "
|
||
|
|
<< (conn->connected() ? "UP" : "DOWN");
|
||
|
|
if (!conn->connected())
|
||
|
|
{
|
||
|
|
if (--senders_ == 0)
|
||
|
|
{
|
||
|
|
output();
|
||
|
|
loop_->quit();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void onMessage(const TcpConnectionPtr& conn, Buffer* buf, Timestamp)
|
||
|
|
{
|
||
|
|
const char* crlf = NULL;
|
||
|
|
while ( (crlf = buf->findCRLF()) != NULL)
|
||
|
|
{
|
||
|
|
// string request(buf->peek(), crlf);
|
||
|
|
// printf("%s\n", request.c_str());
|
||
|
|
const char* tab = std::find(buf->peek(), crlf, '\t');
|
||
|
|
if (tab != crlf)
|
||
|
|
{
|
||
|
|
string word(buf->peek(), tab);
|
||
|
|
int64_t cnt = atoll(tab);
|
||
|
|
wordcounts_[word] += cnt;
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
LOG_ERROR << "Wrong format, no tab found";
|
||
|
|
conn->shutdown();
|
||
|
|
}
|
||
|
|
buf->retrieveUntil(crlf + 2);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void output()
|
||
|
|
{
|
||
|
|
LOG_INFO << "Writing shard";
|
||
|
|
std::ofstream out("shard");
|
||
|
|
for (WordCountMap::iterator it = wordcounts_.begin();
|
||
|
|
it != wordcounts_.end(); ++it)
|
||
|
|
{
|
||
|
|
out << it->first << '\t' << it->second << '\n';
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
EventLoop* loop_;
|
||
|
|
TcpServer server_;
|
||
|
|
int senders_;
|
||
|
|
WordCountMap wordcounts_;
|
||
|
|
};
|
||
|
|
|
||
|
|
int main(int argc, char* argv[])
|
||
|
|
{
|
||
|
|
if (argc < 3)
|
||
|
|
{
|
||
|
|
printf("Usage: %s listen_port number_of_senders\n", argv[0]);
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
EventLoop loop;
|
||
|
|
int port = atoi(argv[1]);
|
||
|
|
InetAddress addr(static_cast<uint16_t>(port));
|
||
|
|
WordCountReceiver receiver(&loop, addr);
|
||
|
|
receiver.start(atoi(argv[2]));
|
||
|
|
loop.loop();
|
||
|
|
}
|
||
|
|
}
|