<commit_before>#include <handy/handy.h>

using namespace std;
using namespace handy;

struct Report {
    long connected;
    long retry;
    long sended;
    long recved;
    Report() { memset(this, 0, sizeof(*this)); }
};

int main(int argc, const char* argv[]) {
    if (argc < 9) {
        printf("usage %s <host> <begin port> <end port> <conn count> <create seconds> <subprocesses> <hearbeat interval> <send size> <management port>\n", argv[0]);
        return 1;
    }
    int c = 1;
    string host = argv[c++];
    int begin_port = atoi(argv[c++]);
    int end_port = atoi(argv[c++]);
    int conn_count = atoi(argv[c++]);
    int create_seconds = atoi(argv[c++]);
    int processes = atoi(argv[c++]);
    conn_count = conn_count / processes;
    int heartbeat_interval = atoi(argv[c++]);
    int bsz = atoi(argv[c++]);
    int man_port = atoi(argv[c++]);

    int pid = 1;
    for (int i = 0; i < processes; i ++) {
        pid = fork();
        if (pid == 0) { // a child process, break
            break;
        }
    }

    EventBase base;
    if (pid == 0) { //child process
        char *buf = new char[bsz];
        ExitCaller ec1([=] {delete[] buf; });
        Slice msg(buf, bsz);
        char heartbeat[] = "heartbeat";
        int send = 0;
        int connected = 0;
        int retry = 0;
        int recved = 0;

        vector<TcpConnPtr> allConns;
        info("creating %d connections", conn_count);
        for (int k = 0; k < create_seconds; k ++) {
            base.runAfter(1000*k, [&]{
                int c = conn_count / create_seconds;
                for (int i = 0; i < c; i++) {
                    short port = begin_port + (i % (end_port - begin_port));
                    auto con = TcpConn::createConnection(&base, host, port, 3000);
                    allConns.push_back(con);
                    con->setReconnectInterval(3000);
                    con->onMsg(new LengthCodec, [&](const TcpConnPtr& con, const Slice& msg) {
                        if (heartbeat_interval == 0) { // echo the msg if no interval
                            con->sendMsg(msg);
                            send++;
                        }
                        recved ++;
                    });
                    con->onState([&, i](const TcpConnPtr &con) {
                        TcpConn::State st = con->getState();
                        if (st == TcpConn::Connected) {
                            connected++;
                            send ++;
                            con->sendMsg(msg);
                        } else if (st == TcpConn::Failed || st == TcpConn::Closed) { //连接出错
                            if (st == TcpConn::Closed) { connected--; }
                            retry++;
                        }
                    });
                }

            });
        }
        if (heartbeat_interval) {
            base.runAfter(1 * 1000, [&] {
                for (size_t i = 0; i < allConns.size(); i++) {
                    if (allConns[i]->getState() == TcpConn::Connected) {
                        allConns[i]->sendMsg(msg);
                        send++;
                    }
                }
            }, heartbeat_interval * 1000);
        }
        TcpConnPtr report = TcpConn::createConnection(&base, "127.0.0.1", man_port, 3000);
        report->onMsg(new LineCodec, [&](const TcpConnPtr& con, Slice msg) {
            if (msg == "exit") {
                info("recv exit msg from master, so exit");
                base.exit();
            }
        });
        report->onState([&](const TcpConnPtr& con) {
            if (con->getState() == TcpConn::Closed) {
                base.exit();
            }
        });
        base.runAfter(100, [&]() {
            report->sendMsg(util::format("%d connected: %ld retry: %ld send: %ld recved: %ld",
                                         getpid(), connected, retry, send, recved));
        }, 100);
        base.loop();
    } else { // master process
        map<int, Report> subs;
        TcpServerPtr master = TcpServer::startServer(&base, "127.0.0.1", man_port);
        master->onConnMsg(new LineCodec, [&](const TcpConnPtr& con, Slice msg) {
            auto fs = msg.split(' ');
            if (fs.size() != 9) {
                error("number of fields is %lu expected 7", fs.size());
                return;
            }
            Report& c = subs[atoi(fs[0].data())];
            c.connected = atoi(fs[2].data());
            c.retry= atoi(fs[4].data());
            c.sended = atoi(fs[6].data());
            c.recved = atoi(fs[8].data());
        });
        base.runAfter(3000, [&](){
            for(auto& s: subs) {
                Report& r = s.second;
                printf("pid: %6ld connected %6ld retry %6ld sended %6ld recved %6ld\n", (long)s.first, r.connected, r.retry, r.sended, r.recved);
            }
            printf("\n");
        }, 3000);
        base.loop();
    }
    info("program exited");
}
<commit_msg>updated from interval send<commit_after>#include <handy/handy.h>

using namespace std;
using namespace handy;

struct Report {
    long connected;
    long retry;
    long sended;
    long recved;
    Report() { memset(this, 0, sizeof(*this)); }
};

int main(int argc, const char* argv[]) {
    if (argc < 9) {
        printf("usage %s <host> <begin port> <end port> <conn count> <create seconds> <subprocesses> <hearbeat interval> <send size> <management port>\n", argv[0]);
        return 1;
    }
    int c = 1;
    string host = argv[c++];
    int begin_port = atoi(argv[c++]);
    int end_port = atoi(argv[c++]);
    int conn_count = atoi(argv[c++]);
    int create_seconds = atoi(argv[c++]);
    int processes = atoi(argv[c++]);
    conn_count = conn_count / processes;
    int heartbeat_interval = atoi(argv[c++]);
    int bsz = atoi(argv[c++]);
    int man_port = atoi(argv[c++]);

    int pid = 1;
    for (int i = 0; i < processes; i ++) {
        pid = fork();
        if (pid == 0) { // a child process, break
            break;
        }
    }

    EventBase base;
    if (pid == 0) { //child process
        char *buf = new char[bsz];
        ExitCaller ec1([=] {delete[] buf; });
        Slice msg(buf, bsz);
        char heartbeat[] = "heartbeat";
        int send = 0;
        int connected = 0;
        int retry = 0;
        int recved = 0;

        vector<TcpConnPtr> allConns;
        info("creating %d connections", conn_count);
        for (int k = 0; k < create_seconds; k ++) {
            base.runAfter(1000*k, [&]{
                int c = conn_count / create_seconds;
                for (int i = 0; i < c; i++) {
                    short port = begin_port + (i % (end_port - begin_port));
                    auto con = TcpConn::createConnection(&base, host, port, 20*000);
                    allConns.push_back(con);
                    con->setReconnectInterval(20*1000);
                    con->onMsg(new LengthCodec, [&](const TcpConnPtr& con, const Slice& msg) {
                        if (heartbeat_interval == 0) { // echo the msg if no interval
                            con->sendMsg(msg);
                            send++;
                        }
                        recved ++;
                    });
                    con->onState([&, i](const TcpConnPtr &con) {
                        TcpConn::State st = con->getState();
                        if (st == TcpConn::Connected) {
                            connected++;
//                            send ++;
//                            con->sendMsg(msg);
                        } else if (st == TcpConn::Failed || st == TcpConn::Closed) { //连接出错
                            if (st == TcpConn::Closed) { connected--; }
                            retry++;
                        }
                    });
                }

            });
        }
        if (heartbeat_interval) {
            base.runAfter(heartbeat_interval * 1000, [&] {
                for (int i = 0; i < heartbeat_interval; i ++) {
                    base.runAfter(i*1000, [&,i]{
                        size_t block = allConns.size() / heartbeat_interval;
                        for (size_t j=i*block; j<(i+1)*block && j<allConns.size(); j++) {
                            if (allConns[i]->getState() == TcpConn::Connected) {
                                allConns[i]->sendMsg(msg);
                                send++;
                            }
                        }
                    });
                }
            }, heartbeat_interval * 1000);
        }
        TcpConnPtr report = TcpConn::createConnection(&base, "127.0.0.1", man_port, 3000);
        report->onMsg(new LineCodec, [&](const TcpConnPtr& con, Slice msg) {
            if (msg == "exit") {
                info("recv exit msg from master, so exit");
                base.exit();
            }
        });
        report->onState([&](const TcpConnPtr& con) {
            if (con->getState() == TcpConn::Closed) {
                base.exit();
            }
        });
        base.runAfter(100, [&]() {
            report->sendMsg(util::format("%d connected: %ld retry: %ld send: %ld recved: %ld",
                                         getpid(), connected, retry, send, recved));
        }, 100);
        base.loop();
    } else { // master process
        map<int, Report> subs;
        TcpServerPtr master = TcpServer::startServer(&base, "127.0.0.1", man_port);
        master->onConnMsg(new LineCodec, [&](const TcpConnPtr& con, Slice msg) {
            auto fs = msg.split(' ');
            if (fs.size() != 9) {
                error("number of fields is %lu expected 7", fs.size());
                return;
            }
            Report& c = subs[atoi(fs[0].data())];
            c.connected = atoi(fs[2].data());
            c.retry= atoi(fs[4].data());
            c.sended = atoi(fs[6].data());
            c.recved = atoi(fs[8].data());
        });
        base.runAfter(3000, [&](){
            for(auto& s: subs) {
                Report& r = s.second;
                printf("pid: %6ld connected %6ld retry %6ld sended %6ld recved %6ld\n", (long)s.first, r.connected, r.retry, r.sended, r.recved);
            }
            printf("\n");
        }, 3000);
        base.loop();
    }
    info("program exited");
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015 OpenALPR Technology, Inc.
 * Open source Automated License Plate Recognition [http://www.openalpr.com]
 *
 * This file is part of OpenALPR.
 *
 * OpenALPR is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License
 * version 3 as published by the Free Software Foundation
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <opencv2/core/core.hpp>

#include "licenseplatecandidate.h"
#include "edges/edgefinder.h"
#include "transformation.h"

using namespace std;
using namespace cv;

namespace alpr
{

  LicensePlateCandidate::LicensePlateCandidate(PipelineData* pipeline_data)
  {
    this->pipeline_data = pipeline_data;
    this->config = pipeline_data->config;

  }

  LicensePlateCandidate::~LicensePlateCandidate()
  {
    delete charSegmenter;
  }

  // Must delete this pointer in parent class
  void LicensePlateCandidate::recognize()
  {
    charSegmenter = NULL;

    pipeline_data->isMultiline = config->multiline;

    Rect expandedRegion = this->pipeline_data->regionOfInterest;

    pipeline_data->crop_gray = Mat(this->pipeline_data->grayImg, expandedRegion);
    resize(pipeline_data->crop_gray, pipeline_data->crop_gray, Size(config->templateWidthPx, config->templateHeightPx));

    CharacterAnalysis textAnalysis(pipeline_data);

    if (pipeline_data->disqualified)
      return;

    EdgeFinder edgeFinder(pipeline_data);

    pipeline_data->plate_corners = edgeFinder.findEdgeCorners();

    if (pipeline_data->disqualified)
      return;

    timespec startTime;
    getTimeMonotonic(&startTime);

    // Compute the transformation matrix to go from the current image to the new plate corners
    Transformation imgTransform(this->pipeline_data->grayImg, pipeline_data->crop_gray, expandedRegion);
    Size cropSize = imgTransform.getCropSize(pipeline_data->plate_corners,
            Size(pipeline_data->config->ocrImageWidthPx, pipeline_data->config->ocrImageHeightPx));
    Mat transmtx = imgTransform.getTransformationMatrix(pipeline_data->plate_corners, cropSize);

    // Crop the plate corners from the original color image (after un-applying prewarp)
    vector<Point2f> projectedPoints = pipeline_data->prewarp->projectPoints(pipeline_data->plate_corners, true);
    pipeline_data->color_deskewed = Mat::zeros(cropSize, pipeline_data->colorImg.type());
    std::vector<cv::Point2f> deskewed_points;
    deskewed_points.push_back(cv::Point2f(0,0));
    deskewed_points.push_back(cv::Point2f(pipeline_data->color_deskewed.cols,0));
    deskewed_points.push_back(cv::Point2f(pipeline_data->color_deskewed.cols,pipeline_data->color_deskewed.rows));
    deskewed_points.push_back(cv::Point2f(0,pipeline_data->color_deskewed.rows));
    cv::Mat color_transmtx = cv::getPerspectiveTransform(projectedPoints, deskewed_points);
    cv::warpPerspective(pipeline_data->colorImg, pipeline_data->color_deskewed, color_transmtx, pipeline_data->color_deskewed.size());

    if (pipeline_data->color_deskewed.channels() > 2)
    {
      // Make a grayscale copy as well for faster processing downstream
      cv::cvtColor(pipeline_data->color_deskewed, pipeline_data->crop_gray, CV_BGR2GRAY);
    }

    if (this->config->debugGeneral)
      displayImage(config, "quadrilateral", pipeline_data->color_deskewed);

    // Apply a perspective transformation to the TextLine objects
    // to match the newly deskewed license plate crop
    vector<TextLine> newLines;
    for (unsigned int i = 0; i < pipeline_data->textLines.size(); i++)
    {
      vector<Point2f> textArea = imgTransform.transformSmallPointsToBigImage(pipeline_data->textLines[i].textArea);
      vector<Point2f> linePolygon = imgTransform.transformSmallPointsToBigImage(pipeline_data->textLines[i].linePolygon);

      vector<Point2f> textAreaRemapped;
      vector<Point2f> linePolygonRemapped;

      textAreaRemapped = imgTransform.remapSmallPointstoCrop(textArea, transmtx);
      linePolygonRemapped = imgTransform.remapSmallPointstoCrop(linePolygon, transmtx);

      newLines.push_back(TextLine(textAreaRemapped, linePolygonRemapped, pipeline_data->crop_gray.size()));
    }

    pipeline_data->textLines.clear();
    for (unsigned int i = 0; i < newLines.size(); i++)
      pipeline_data->textLines.push_back(newLines[i]);

    if (config->debugTiming)
    {
      timespec endTime;
      getTimeMonotonic(&endTime);
      cout << "deskew Time: " << diffclock(startTime, endTime) << "ms." << endl;
    }

    charSegmenter = new CharacterSegmenter(pipeline_data);

  }

}
<commit_msg>Added copy to grayscale img crop related to issue #247 and #216<commit_after>/*
 * Copyright (c) 2015 OpenALPR Technology, Inc.
 * Open source Automated License Plate Recognition [http://www.openalpr.com]
 *
 * This file is part of OpenALPR.
 *
 * OpenALPR is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License
 * version 3 as published by the Free Software Foundation
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <opencv2/core/core.hpp>

#include "licenseplatecandidate.h"
#include "edges/edgefinder.h"
#include "transformation.h"

using namespace std;
using namespace cv;

namespace alpr
{

  LicensePlateCandidate::LicensePlateCandidate(PipelineData* pipeline_data)
  {
    this->pipeline_data = pipeline_data;
    this->config = pipeline_data->config;

  }

  LicensePlateCandidate::~LicensePlateCandidate()
  {
    delete charSegmenter;
  }

  // Must delete this pointer in parent class
  void LicensePlateCandidate::recognize()
  {
    charSegmenter = NULL;

    pipeline_data->isMultiline = config->multiline;

    Rect expandedRegion = this->pipeline_data->regionOfInterest;

    pipeline_data->crop_gray = Mat(this->pipeline_data->grayImg, expandedRegion);
    resize(pipeline_data->crop_gray, pipeline_data->crop_gray, Size(config->templateWidthPx, config->templateHeightPx));

    CharacterAnalysis textAnalysis(pipeline_data);

    if (pipeline_data->disqualified)
      return;

    EdgeFinder edgeFinder(pipeline_data);

    pipeline_data->plate_corners = edgeFinder.findEdgeCorners();

    if (pipeline_data->disqualified)
      return;

    timespec startTime;
    getTimeMonotonic(&startTime);

    // Compute the transformation matrix to go from the current image to the new plate corners
    Transformation imgTransform(this->pipeline_data->grayImg, pipeline_data->crop_gray, expandedRegion);
    Size cropSize = imgTransform.getCropSize(pipeline_data->plate_corners,
            Size(pipeline_data->config->ocrImageWidthPx, pipeline_data->config->ocrImageHeightPx));
    Mat transmtx = imgTransform.getTransformationMatrix(pipeline_data->plate_corners, cropSize);

    // Crop the plate corners from the original color image (after un-applying prewarp)
    vector<Point2f> projectedPoints = pipeline_data->prewarp->projectPoints(pipeline_data->plate_corners, true);
    pipeline_data->color_deskewed = Mat::zeros(cropSize, pipeline_data->colorImg.type());
    std::vector<cv::Point2f> deskewed_points;
    deskewed_points.push_back(cv::Point2f(0,0));
    deskewed_points.push_back(cv::Point2f(pipeline_data->color_deskewed.cols,0));
    deskewed_points.push_back(cv::Point2f(pipeline_data->color_deskewed.cols,pipeline_data->color_deskewed.rows));
    deskewed_points.push_back(cv::Point2f(0,pipeline_data->color_deskewed.rows));
    cv::Mat color_transmtx = cv::getPerspectiveTransform(projectedPoints, deskewed_points);
    cv::warpPerspective(pipeline_data->colorImg, pipeline_data->color_deskewed, color_transmtx, pipeline_data->color_deskewed.size());

    if (pipeline_data->color_deskewed.channels() > 2)
    {
      // Make a grayscale copy as well for faster processing downstream
      cv::cvtColor(pipeline_data->color_deskewed, pipeline_data->crop_gray, CV_BGR2GRAY);
    }
    else
    {
      // Copy the already grayscale image to the crop_gray img
      pipeline_data->color_deskewed.copyTo(pipeline_data->crop_gray);
    }

    if (this->config->debugGeneral)
      displayImage(config, "quadrilateral", pipeline_data->color_deskewed);

    // Apply a perspective transformation to the TextLine objects
    // to match the newly deskewed license plate crop
    vector<TextLine> newLines;
    for (unsigned int i = 0; i < pipeline_data->textLines.size(); i++)
    {
      vector<Point2f> textArea = imgTransform.transformSmallPointsToBigImage(pipeline_data->textLines[i].textArea);
      vector<Point2f> linePolygon = imgTransform.transformSmallPointsToBigImage(pipeline_data->textLines[i].linePolygon);

      vector<Point2f> textAreaRemapped;
      vector<Point2f> linePolygonRemapped;

      textAreaRemapped = imgTransform.remapSmallPointstoCrop(textArea, transmtx);
      linePolygonRemapped = imgTransform.remapSmallPointstoCrop(linePolygon, transmtx);

      newLines.push_back(TextLine(textAreaRemapped, linePolygonRemapped, pipeline_data->crop_gray.size()));
    }

    pipeline_data->textLines.clear();
    for (unsigned int i = 0; i < newLines.size(); i++)
      pipeline_data->textLines.push_back(newLines[i]);

    if (config->debugTiming)
    {
      timespec endTime;
      getTimeMonotonic(&endTime);
      cout << "deskew Time: " << diffclock(startTime, endTime) << "ms." << endl;
    }

    charSegmenter = new CharacterSegmenter(pipeline_data);

  }

}
<|endoftext|>


<commit_before>#ifndef _X_EVENT_HANDLER
#define _X_EVENT_HANDLER

class x_event_handler {
  public:
    virtual void handle(xcb_generic_event_t *) = 0;
};

#endif
<commit_msg>Add missing include<commit_after>#ifndef _X_EVENT_HANDLER
#define _X_EVENT_HANDLER

#include <xcb/xcb.h>

class x_event_handler {
  public:
    virtual void handle(xcb_generic_event_t *) = 0;
};

#endif
<|endoftext|>


<commit_before>#include <vast/program.h>

#include <cstdlib>
#include <iostream>
#include <boost/exception/diagnostic_information.hpp>
#include <vast/exception.h>
#include <vast/comm/broccoli.h>
#include <vast/detail/cppa_type_info.h>
#include <vast/fs/path.h>
#include <vast/fs/operations.h>
#include <vast/ingest/ingestor.h>
#include <vast/meta/schema_manager.h>
#include <vast/query/client.h>
#include <vast/query/search.h>
#include <vast/store/archive.h>
#include <vast/store/index.h>
#include <vast/util/logger.h>
#include <config.h>

#ifdef USE_PERFTOOLS_CPU_PROFILER
#include <google/profiler.h>
#endif
#ifdef USE_PERFTOOLS_HEAP_PROFILER
#include <google/heap-profiler.h>
#endif

namespace vast {

/// Declaration of global (extern) variables.
namespace util {
logger* LOGGER;
}

program::program()
  : terminating_(false)
  , return_(EXIT_FAILURE)
{
}

program::~program()
{
}

bool program::init(std::string const& filename)
{
  try
  {
    config_.load(filename);
    do_init();
    return true;
  }
  catch (boost::program_options::unknown_option const& e)
  {
    std::cerr << e.what();
  }

  return false;
}

bool program::init(int argc, char *argv[])
{
  try
  {
    config_.load(argc, argv);

    if (argc < 2 || config_.check("help") || config_.check("advanced"))
    {
      config_.print(std::cerr, config_.check("advanced"));
      return false;
    }

    do_init();
    return true;
  }
  catch (config_exception const& e)
  {
    std::cerr << e.what() << std::endl;
  }
  catch (boost::program_options::unknown_option const& e)
  {
    std::cerr << e.what() << ", try -h or --help" << std::endl;
  }
  catch (boost::program_options::invalid_command_line_syntax const& e)
  {
    std::cerr << "invalid command line: " << e.what() << std::endl;
  }
  catch (boost::exception const& e)
  {
    std::cerr << boost::diagnostic_information(e);
  }

  return false;
}

void program::start()
{
  using namespace cppa;
  detail::cppa_announce_types();

  auto log_dir = config_.get<fs::path>("directory") / "log";
  assert(fs::exists(log_dir));

  try
  {
#ifdef USE_PERFTOOLS_HEAP_PROFILER
    if (config_.check("perftools-heap"))
    {
      LOG(info, core) << "starting Gperftools CPU profiler";
      HeapProfilerStart((log_dir / "heap.profile").string().data());
    }
#endif
#ifdef USE_PERFTOOLS_CPU_PROFILER
    if (config_.check("perftools-cpu"))
    {
      LOG(info, core) << "starting Gperftools heap profiler";
      ProfilerStart((log_dir / "cpu.profile").string().data());
    }
#endif

    if (config_.check("profile"))
    {
      auto const& filename = log_dir / "profiler.log";
      auto ms = config_.get<unsigned>("profile");
      profiler_ = spawn<util::profiler>(filename.string(),
                                        std::chrono::seconds(ms));
      send(profiler_, atom("run"));
    }

    comm::broccoli::init(config_.check("broccoli-messages"),
                         config_.check("broccoli-calltrace"));

    schema_manager_ = spawn<meta::schema_manager>();
    if (config_.check("schema"))
    {
      send(schema_manager_, atom("load"), config_.get<std::string>("schema"));

      if (config_.check("print-schema"))
      {
        send(schema_manager_, atom("print"));
        receive(
            on(atom("schema"), arg_match) >> [](std::string const& schema)
            {
              std::cout << schema << std::endl;
            });

        return;
      }
    }

    if (config_.check("archive-actor"))
      archive_ = spawn<store::archive>(
          (config_.get<fs::path>("directory") / "archive").string(),
          config_.get<size_t>("archive.max-events-per-chunk"),
          config_.get<size_t>("archive.max-segment-size") * 1000,
          config_.get<size_t>("archive.max-segments"));

    if (config_.check("index-actor"))
        index_ = spawn<store::index>(
            archive_,
            (config_.get<fs::path>("directory") / "index").string());

    if (config_.check("ingestor-actor"))
    {
      ingestor_ = spawn<ingest::ingestor>(archive_);
      send(ingestor_,
           atom("initialize"),
           config_.get<std::string>("ingestor.host"),
           config_.get<unsigned>("ingestor.port"));

      if (config_.check("ingestor.events"))
      {
        auto events = config_.get<std::vector<std::string>>("ingestor.events");
        for (auto& event : events)
          send(ingestor_, atom("subscribe"), event);
      }

      if (config_.check("ingestor.file"))
      {
        auto files = config_.get<std::vector<std::string>>("ingestor.file");
        for (auto& file : files)
        {
          LOG(info, core) << "ingesting " << file;
          send(ingestor_, atom("read_file"), file);
        }
      }
    }

    if (config_.check("search-actor"))
    {
      search_ = spawn<query::search>(archive_, index_);

      //config_.get<std::string>("search.host")
      LOG(verbose, core) << "publishing search at *:"
          << config_.get<unsigned>("search.port");
      publish(search_, config_.get<unsigned>("search.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to search at "
          << config_.get<std::string>("search.host") << ":"
          << config_.get<unsigned>("search.port");
      search_ = remote_actor(
          config_.get<std::string>("search.host"),
          config_.get<unsigned>("search.port"));
    }

    if (config_.check("query"))
    {
      auto paginate = config_.get<unsigned>("client.paginate");
      auto& expression = config_.get<std::string>("query");
      query_client_ = spawn<query::client>(search_, paginate);
      send(query_client_, atom("query"), atom("create"), expression);
    }

    await_all_others_done();
    return_ = EXIT_SUCCESS;
  }
  catch (network_error const& e)
  {
      LOG(error, core) << "network error: " << e.what();
  }
  catch (...)
  {
    LOG(fatal, core)
      << "exception details:\n"
      << boost::current_exception_diagnostic_information();
  }
}

void program::stop()
{
  using namespace cppa;

  if (terminating_)
  {
    return_ = EXIT_FAILURE;
    return;
  }

  terminating_ = true;

  auto shutdown = make_any_tuple(atom("shutdown"));

  if (config_.check("query"))
    query_client_ << shutdown;

  if (config_.check("search-actor"))
    search_ << shutdown;

  if (config_.check("ingestor-actor"))
    ingestor_ << shutdown;

  if (config_.check("index-actor"))
    index_ << shutdown;

  if (config_.check("archive-actor"))
    archive_ << shutdown;

  schema_manager_ << shutdown;

  if (config_.check("profile"))
    profiler_ << shutdown;

#ifdef USE_PERFTOOLS_CPU_PROFILER
  ProfilerState state;
  ProfilerGetCurrentState(&state);
  LOG(info, core)
    << "Gperftools CPU profiler gathered "
    <<  state.samples_gathered << " samples"
    << " in file " << state.profile_name;

  if (config_.check("perftools-cpu"))
  {
    LOG(info, core) << "stopping Gperftools CPU profiler";
    ProfilerStop();
  }
#endif
#ifdef USE_PERFTOOLS_HEAP_PROFILER
  if (config_.check("perftools-heap") && IsHeapProfilerRunning())
  {
    LOG(info, core) << "stopping Gperftools heap profiler";
    HeapProfilerDump("cleanup");
    HeapProfilerStop();
  }
#endif

  return_ = EXIT_SUCCESS;
}

int program::end()
{
  switch (return_)
  {
    case EXIT_SUCCESS:
      LOG(info, core) << "vast terminated cleanly";
      break;

    case EXIT_FAILURE:
      LOG(info, core) << "vast terminated with errors";
      break;

    default:
      assert(! "invalid return code");
  }

  return return_;
}

void program::do_init()
{
  auto vast_dir = config_.get<fs::path>("directory");
  if (! fs::exists(vast_dir))
    fs::mkdir(vast_dir);

  auto log_dir = vast_dir / "log";
  if (! fs::exists(log_dir))
      fs::mkdir(log_dir);

  util::LOGGER = new util::logger(
      static_cast<util::logger::level>(config_.get<int>("console-verbosity")),
      static_cast<util::logger::level>(config_.get<int>("logfile-verbosity")),
      log_dir / "vast.log");

  LOG(verbose, core) << " _   _____   __________";
  LOG(verbose, core) << "| | / / _ | / __/_  __/";
  LOG(verbose, core) << "| |/ / __ |_\\ \\  / / ";
  LOG(verbose, core) << "|___/_/ |_/___/ /_/  " << VAST_VERSION;
  LOG(verbose, core) << "";
}

} // namespace vast
<commit_msg>Fix Gperftools logging output.<commit_after>#include <vast/program.h>

#include <cstdlib>
#include <iostream>
#include <boost/exception/diagnostic_information.hpp>
#include <vast/exception.h>
#include <vast/comm/broccoli.h>
#include <vast/detail/cppa_type_info.h>
#include <vast/fs/path.h>
#include <vast/fs/operations.h>
#include <vast/ingest/ingestor.h>
#include <vast/meta/schema_manager.h>
#include <vast/query/client.h>
#include <vast/query/search.h>
#include <vast/store/archive.h>
#include <vast/store/index.h>
#include <vast/util/logger.h>
#include <config.h>

#ifdef USE_PERFTOOLS_CPU_PROFILER
#include <google/profiler.h>
#endif
#ifdef USE_PERFTOOLS_HEAP_PROFILER
#include <google/heap-profiler.h>
#endif

namespace vast {

/// Declaration of global (extern) variables.
namespace util {
logger* LOGGER;
}

program::program()
  : terminating_(false)
  , return_(EXIT_FAILURE)
{
}

program::~program()
{
}

bool program::init(std::string const& filename)
{
  try
  {
    config_.load(filename);
    do_init();
    return true;
  }
  catch (boost::program_options::unknown_option const& e)
  {
    std::cerr << e.what();
  }

  return false;
}

bool program::init(int argc, char *argv[])
{
  try
  {
    config_.load(argc, argv);

    if (argc < 2 || config_.check("help") || config_.check("advanced"))
    {
      config_.print(std::cerr, config_.check("advanced"));
      return false;
    }

    do_init();
    return true;
  }
  catch (config_exception const& e)
  {
    std::cerr << e.what() << std::endl;
  }
  catch (boost::program_options::unknown_option const& e)
  {
    std::cerr << e.what() << ", try -h or --help" << std::endl;
  }
  catch (boost::program_options::invalid_command_line_syntax const& e)
  {
    std::cerr << "invalid command line: " << e.what() << std::endl;
  }
  catch (boost::exception const& e)
  {
    std::cerr << boost::diagnostic_information(e);
  }

  return false;
}

void program::start()
{
  using namespace cppa;
  detail::cppa_announce_types();

  auto log_dir = config_.get<fs::path>("directory") / "log";
  assert(fs::exists(log_dir));

  try
  {
#ifdef USE_PERFTOOLS_HEAP_PROFILER
    if (config_.check("perftools-heap"))
    {
      LOG(info, core) << "starting Gperftools CPU profiler";
      HeapProfilerStart((log_dir / "heap.profile").string().data());
    }
#endif
#ifdef USE_PERFTOOLS_CPU_PROFILER
    if (config_.check("perftools-cpu"))
    {
      LOG(info, core) << "starting Gperftools heap profiler";
      ProfilerStart((log_dir / "cpu.profile").string().data());
    }
#endif

    if (config_.check("profile"))
    {
      auto const& filename = log_dir / "profiler.log";
      auto ms = config_.get<unsigned>("profile");
      profiler_ = spawn<util::profiler>(filename.string(),
                                        std::chrono::seconds(ms));
      send(profiler_, atom("run"));
    }

    comm::broccoli::init(config_.check("broccoli-messages"),
                         config_.check("broccoli-calltrace"));

    schema_manager_ = spawn<meta::schema_manager>();
    if (config_.check("schema"))
    {
      send(schema_manager_, atom("load"), config_.get<std::string>("schema"));

      if (config_.check("print-schema"))
      {
        send(schema_manager_, atom("print"));
        receive(
            on(atom("schema"), arg_match) >> [](std::string const& schema)
            {
              std::cout << schema << std::endl;
            });

        return;
      }
    }

    if (config_.check("archive-actor"))
      archive_ = spawn<store::archive>(
          (config_.get<fs::path>("directory") / "archive").string(),
          config_.get<size_t>("archive.max-events-per-chunk"),
          config_.get<size_t>("archive.max-segment-size") * 1000,
          config_.get<size_t>("archive.max-segments"));

    if (config_.check("index-actor"))
        index_ = spawn<store::index>(
            archive_,
            (config_.get<fs::path>("directory") / "index").string());

    if (config_.check("ingestor-actor"))
    {
      ingestor_ = spawn<ingest::ingestor>(archive_);
      send(ingestor_,
           atom("initialize"),
           config_.get<std::string>("ingestor.host"),
           config_.get<unsigned>("ingestor.port"));

      if (config_.check("ingestor.events"))
      {
        auto events = config_.get<std::vector<std::string>>("ingestor.events");
        for (auto& event : events)
          send(ingestor_, atom("subscribe"), event);
      }

      if (config_.check("ingestor.file"))
      {
        auto files = config_.get<std::vector<std::string>>("ingestor.file");
        for (auto& file : files)
        {
          LOG(info, core) << "ingesting " << file;
          send(ingestor_, atom("read_file"), file);
        }
      }
    }

    if (config_.check("search-actor"))
    {
      search_ = spawn<query::search>(archive_, index_);

      //config_.get<std::string>("search.host")
      LOG(verbose, core) << "publishing search at *:"
          << config_.get<unsigned>("search.port");
      publish(search_, config_.get<unsigned>("search.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to search at "
          << config_.get<std::string>("search.host") << ":"
          << config_.get<unsigned>("search.port");
      search_ = remote_actor(
          config_.get<std::string>("search.host"),
          config_.get<unsigned>("search.port"));
    }

    if (config_.check("query"))
    {
      auto paginate = config_.get<unsigned>("client.paginate");
      auto& expression = config_.get<std::string>("query");
      query_client_ = spawn<query::client>(search_, paginate);
      send(query_client_, atom("query"), atom("create"), expression);
    }

    await_all_others_done();
    return_ = EXIT_SUCCESS;
  }
  catch (network_error const& e)
  {
      LOG(error, core) << "network error: " << e.what();
  }
  catch (...)
  {
    LOG(fatal, core)
      << "exception details:\n"
      << boost::current_exception_diagnostic_information();
  }
}

void program::stop()
{
  using namespace cppa;

  if (terminating_)
  {
    return_ = EXIT_FAILURE;
    return;
  }

  terminating_ = true;

  auto shutdown = make_any_tuple(atom("shutdown"));

  if (config_.check("query"))
    query_client_ << shutdown;

  if (config_.check("search-actor"))
    search_ << shutdown;

  if (config_.check("ingestor-actor"))
    ingestor_ << shutdown;

  if (config_.check("index-actor"))
    index_ << shutdown;

  if (config_.check("archive-actor"))
    archive_ << shutdown;

  schema_manager_ << shutdown;

  if (config_.check("profile"))
    profiler_ << shutdown;

#ifdef USE_PERFTOOLS_CPU_PROFILER

  if (config_.check("perftools-cpu"))
  {
    ProfilerState state;
    ProfilerGetCurrentState(&state);
    LOG(info, core)
      << "Gperftools CPU profiler gathered "
      <<  state.samples_gathered << " samples"
      << " in file " << state.profile_name;

    LOG(info, core) << "stopping Gperftools CPU profiler";
    ProfilerStop();
  }
#endif
#ifdef USE_PERFTOOLS_HEAP_PROFILER
  if (config_.check("perftools-heap") && IsHeapProfilerRunning())
  {
    LOG(info, core) << "stopping Gperftools heap profiler";
    HeapProfilerDump("cleanup");
    HeapProfilerStop();
  }
#endif

  return_ = EXIT_SUCCESS;
}

int program::end()
{
  switch (return_)
  {
    case EXIT_SUCCESS:
      LOG(info, core) << "vast terminated cleanly";
      break;

    case EXIT_FAILURE:
      LOG(info, core) << "vast terminated with errors";
      break;

    default:
      assert(! "invalid return code");
  }

  return return_;
}

void program::do_init()
{
  auto vast_dir = config_.get<fs::path>("directory");
  if (! fs::exists(vast_dir))
    fs::mkdir(vast_dir);

  auto log_dir = vast_dir / "log";
  if (! fs::exists(log_dir))
      fs::mkdir(log_dir);

  util::LOGGER = new util::logger(
      static_cast<util::logger::level>(config_.get<int>("console-verbosity")),
      static_cast<util::logger::level>(config_.get<int>("logfile-verbosity")),
      log_dir / "vast.log");

  LOG(verbose, core) << " _   _____   __________";
  LOG(verbose, core) << "| | / / _ | / __/_  __/";
  LOG(verbose, core) << "| |/ / __ |_\\ \\  / / ";
  LOG(verbose, core) << "|___/_/ |_/___/ /_/  " << VAST_VERSION;
  LOG(verbose, core) << "";
}

} // namespace vast
<|endoftext|>


<commit_before>#include "vast/program.h"

#include <cstdlib>
#include <iostream>
#include "vast/archive.h"
#include "vast/exception.h"
#include "vast/id_tracker.h"
#include "vast/index.h"
#include "vast/ingestor.h"
#include "vast/logger.h"
#include "vast/query_client.h"
#include "vast/search.h"
#include "vast/system_monitor.h"
#include "vast/comm/broccoli.h"
#include "vast/detail/cppa_type_info.h"
#include "vast/fs/path.h"
#include "vast/fs/operations.h"
#include "vast/meta/schema_manager.h"
#include "vast/util/profiler.h"
#include "config.h"

namespace vast {

using namespace cppa;

program::program(configuration const& config)
  : config_(config)
{
  detail::cppa_announce_types();
}

bool program::run()
{
  if (! start())
    return false;

  bool done = false;
  do_receive(
      on(atom("system"), atom("keystroke"), arg_match) >> [=](char key)
      {
        if (! query_client_)
          return;

        switch (key)
        {
          default:
            LOG(info, core) << "invalid key: '" << key << "'";
          case '?':
            LOG(info, core)
              << "available commands: "
                 "<space> for results, (s)tatistics, (Q)uit";
            break;
          case ' ':
            send(query_client_, atom("client"), atom("results"));
            break;
          case 's':
            send(query_client_, atom("client"), atom("statistics"));
            break;
        }
      },
      on(atom("DOWN"), arg_match) >> [&done](uint32_t reason)
      {
        done = true;
      }).until(gref(done) == true);

  stop();
  await_all_others_done();

  return true;
}

bool program::start()
{
  LOG(verbose, core) << " _   _____   __________";
  LOG(verbose, core) << "| | / / _ | / __/_  __/";
  LOG(verbose, core) << "| |/ / __ |_\\ \\  / / ";
  LOG(verbose, core) << "|___/_/ |_/___/ /_/  " << VAST_VERSION;
  LOG(verbose, core) << "";

  auto vast_dir = config_.get<fs::path>("directory");
  if (! fs::exists(vast_dir))
    fs::mkdir(vast_dir);

  auto log_dir = config_.get<fs::path>("log.directory");
  if (! fs::exists(log_dir))
      fs::mkdir(log_dir);

  try
  {
    system_monitor_ = spawn<system_monitor>(self);
    self->monitor(system_monitor_);

    if (config_.check("profile"))
    {
      auto ms = config_.get<unsigned>("profile");
      profiler_ = spawn<util::profiler>(log_dir.string(), std::chrono::seconds(ms));
      send(profiler_,
           atom("run"),
           config_.check("profile-cpu"),
           config_.check("profile-heap"));
    }

    // TODO: uncomment once brought back into the game.
    //schema_manager_ = spawn<meta::schema_manager>();
    //if (config_.check("schema"))
    //{
    //  send(schema_manager_, atom("load"), config_.get<std::string>("schema"));

    //  if (config_.check("print-schema"))
    //  {
    //    send(schema_manager_, atom("print"));
    //    receive(
    //        on(atom("schema"), arg_match) >> [](std::string const& schema)
    //        {
    //          std::cout << schema << std::endl;
    //        },
    //        after(std::chrono::seconds(1)) >> [=]
    //        {
    //          LOG(error, meta)
    //            << "schema manager did not answer after one second";
    //        });

    //    return false;
    //  }
    //}

    if (config_.check("tracker-actor") || config_.check("all-server"))
    {
      tracker_ = spawn<id_tracker>(
          (config_.get<fs::path>("directory") / "id").string());

      LOG(verbose, core) << "publishing tracker at *:"
          << config_.get<unsigned>("tracker.port");

      publish(tracker_, config_.get<unsigned>("tracker.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to tracker at "
          << config_.get<std::string>("tracker.host") << ":"
          << config_.get<unsigned>("tracker.port");

      tracker_ = remote_actor(
          config_.get<std::string>("tracker.host"),
          config_.get<unsigned>("tracker.port"));

      LOG(verbose, core) << "connected to tracker actor @" << tracker_->id();
    }

    if (config_.check("archive-actor") || config_.check("all-server"))
    {
      archive_ = spawn<archive>(
          (config_.get<fs::path>("directory") / "archive").string(),
          config_.get<size_t>("archive.max-segments"));
      send(archive_, atom("load"));

      LOG(verbose, core) << "publishing archive at *:"
          << config_.get<unsigned>("archive.port");

      publish(archive_, config_.get<unsigned>("archive.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to archive at "
          << config_.get<std::string>("archive.host") << ":"
          << config_.get<unsigned>("archive.port");

      archive_ = remote_actor(
          config_.get<std::string>("archive.host"),
          config_.get<unsigned>("archive.port"));

      LOG(verbose, core) << "connected to archive actor @" << archive_->id();
    }

    if (config_.check("index-actor") || config_.check("all-server"))
    {
      index_ = spawn<index>(
          archive_,
          (config_.get<fs::path>("directory") / "index").string());
      send(index_, atom("load"));

      LOG(verbose, core) << "publishing index at *:"
          << config_.get<unsigned>("index.port");

      publish(index_, config_.get<unsigned>("index.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to index at "
          << config_.get<std::string>("index.host") << ":"
          << config_.get<unsigned>("index.port");

      index_ = remote_actor(
          config_.get<std::string>("index.host"),
          config_.get<unsigned>("index.port"));

      LOG(verbose, core) << "connected to index actor @" << index_->id();
    }

    if (config_.check("ingestor-actor"))
    {
      comm::broccoli::init(config_.check("broccoli-messages"),
                           config_.check("broccoli-calltrace"));

      ingestor_ = spawn<ingestor>(tracker_, archive_, index_);
      self->monitor(ingestor_);

      send(ingestor_, atom("initialize"),
          config_.get<size_t>("ingest.max-events-per-chunk"),
          config_.get<size_t>("ingest.max-segment-size") * 1000000,
          config_.get<size_t>("ingest.batch-size"));

      if (config_.check("ingest.events"))
      {
        auto host = config_.get<std::string>("ingest.broccoli-host");
        auto port = config_.get<unsigned>("ingest.broccoli-port");
        auto events = config_.get<std::vector<std::string>>("ingest.broccoli-events");
        send(ingestor_, atom("ingest"), atom("broccoli"), host, port, events);
      }

      if (config_.check("ingest.file-names"))
      {
        auto type = config_.get<std::string>("ingest.file-type");
        auto files = config_.get<std::vector<std::string>>("ingest.file-names");
        for (auto& file : files)
        {
          if (fs::exists(file))
            send(ingestor_, atom("ingest"), type, file);
          else
            LOG(error, core) << "no such file: " << file;
        }
      }

      send(ingestor_, atom("extract"));
    }

    if (config_.check("search-actor") || config_.check("all-server"))
    {
      search_ = spawn<search>(archive_, index_);

      LOG(verbose, core) << "publishing search at *:"
          << config_.get<unsigned>("search.port");

      publish(search_, config_.get<unsigned>("search.port"));
    }
    else if (config_.check("client.expression"))
    {
      LOG(verbose, core) << "connecting to search at "
          << config_.get<std::string>("search.host") << ":"
          << config_.get<unsigned>("search.port");

      search_ = remote_actor(
          config_.get<std::string>("search.host"),
          config_.get<unsigned>("search.port"));

      LOG(verbose, core) << "connected to search actor @" << search_->id();

      auto paginate = config_.get<unsigned>("client.paginate");
      auto& expression = config_.get<std::string>("client.expression");
      query_client_ = spawn<query_client>(search_, expression, paginate);
      send(query_client_, atom("start"));
    }

    return true;
  }
  catch (network_error const& e)
  {
      LOG(error, core) << "network error: " << e.what();
  }

  return false;
}

void program::stop()
{
  auto shutdown = make_any_tuple(atom("shutdown"));

  if (query_client_)
    query_client_ << shutdown;

  if (config_.check("search-actor") || config_.check("all-server"))
    search_ << shutdown;

  if (config_.check("ingestor-actor"))
    ingestor_ << shutdown;

  if (config_.check("index-actor") || config_.check("all-server"))
    index_ << shutdown;

  if (config_.check("archive-actor") || config_.check("all-server"))
    archive_ << shutdown;

  if (config_.check("tracker-actor") || config_.check("all-server"))
    tracker_ << shutdown;

  // TODO: uncomment once brought back into the game.
  //schema_manager_ << shutdown;

  if (profiler_)
    profiler_ << shutdown;

  if (system_monitor_)
    system_monitor_ << shutdown;

}

} // namespace vast
<commit_msg>Terminate client when query finishes.<commit_after>#include "vast/program.h"

#include <cstdlib>
#include <iostream>
#include "vast/archive.h"
#include "vast/exception.h"
#include "vast/id_tracker.h"
#include "vast/index.h"
#include "vast/ingestor.h"
#include "vast/logger.h"
#include "vast/query_client.h"
#include "vast/search.h"
#include "vast/system_monitor.h"
#include "vast/comm/broccoli.h"
#include "vast/detail/cppa_type_info.h"
#include "vast/fs/path.h"
#include "vast/fs/operations.h"
#include "vast/meta/schema_manager.h"
#include "vast/util/profiler.h"
#include "config.h"

namespace vast {

using namespace cppa;

program::program(configuration const& config)
  : config_(config)
{
  detail::cppa_announce_types();
}

bool program::run()
{
  if (! start())
    return false;

  bool done = false;
  do_receive(
      on(atom("system"), atom("keystroke"), arg_match) >> [=](char key)
      {
        if (! query_client_)
          return;

        switch (key)
        {
          default:
            LOG(info, core) << "invalid key: '" << key << "'";
          case '?':
            LOG(info, core)
              << "available commands: "
                 "<space> for results, (s)tatistics, (Q)uit";
            break;
          case ' ':
            send(query_client_, atom("client"), atom("results"));
            break;
          case 's':
            send(query_client_, atom("client"), atom("statistics"));
            break;
        }
      },
      on(atom("DOWN"), arg_match) >> [&done](uint32_t reason)
      {
        done = true;
      }).until(gref(done) == true);

  stop();
  await_all_others_done();

  return true;
}

bool program::start()
{
  LOG(verbose, core) << " _   _____   __________";
  LOG(verbose, core) << "| | / / _ | / __/_  __/";
  LOG(verbose, core) << "| |/ / __ |_\\ \\  / / ";
  LOG(verbose, core) << "|___/_/ |_/___/ /_/  " << VAST_VERSION;
  LOG(verbose, core) << "";

  auto vast_dir = config_.get<fs::path>("directory");
  if (! fs::exists(vast_dir))
    fs::mkdir(vast_dir);

  auto log_dir = config_.get<fs::path>("log.directory");
  if (! fs::exists(log_dir))
      fs::mkdir(log_dir);

  try
  {
    system_monitor_ = spawn<system_monitor>(self);
    self->monitor(system_monitor_);

    if (config_.check("profile"))
    {
      auto ms = config_.get<unsigned>("profile");
      profiler_ = spawn<util::profiler>(log_dir.string(), std::chrono::seconds(ms));
      send(profiler_,
           atom("run"),
           config_.check("profile-cpu"),
           config_.check("profile-heap"));
    }

    // TODO: uncomment once brought back into the game.
    //schema_manager_ = spawn<meta::schema_manager>();
    //if (config_.check("schema"))
    //{
    //  send(schema_manager_, atom("load"), config_.get<std::string>("schema"));

    //  if (config_.check("print-schema"))
    //  {
    //    send(schema_manager_, atom("print"));
    //    receive(
    //        on(atom("schema"), arg_match) >> [](std::string const& schema)
    //        {
    //          std::cout << schema << std::endl;
    //        },
    //        after(std::chrono::seconds(1)) >> [=]
    //        {
    //          LOG(error, meta)
    //            << "schema manager did not answer after one second";
    //        });

    //    return false;
    //  }
    //}

    if (config_.check("tracker-actor") || config_.check("all-server"))
    {
      tracker_ = spawn<id_tracker>(
          (config_.get<fs::path>("directory") / "id").string());

      LOG(verbose, core) << "publishing tracker at *:"
          << config_.get<unsigned>("tracker.port");

      publish(tracker_, config_.get<unsigned>("tracker.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to tracker at "
          << config_.get<std::string>("tracker.host") << ":"
          << config_.get<unsigned>("tracker.port");

      tracker_ = remote_actor(
          config_.get<std::string>("tracker.host"),
          config_.get<unsigned>("tracker.port"));

      LOG(verbose, core) << "connected to tracker actor @" << tracker_->id();
    }

    if (config_.check("archive-actor") || config_.check("all-server"))
    {
      archive_ = spawn<archive>(
          (config_.get<fs::path>("directory") / "archive").string(),
          config_.get<size_t>("archive.max-segments"));
      send(archive_, atom("load"));

      LOG(verbose, core) << "publishing archive at *:"
          << config_.get<unsigned>("archive.port");

      publish(archive_, config_.get<unsigned>("archive.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to archive at "
          << config_.get<std::string>("archive.host") << ":"
          << config_.get<unsigned>("archive.port");

      archive_ = remote_actor(
          config_.get<std::string>("archive.host"),
          config_.get<unsigned>("archive.port"));

      LOG(verbose, core) << "connected to archive actor @" << archive_->id();
    }

    if (config_.check("index-actor") || config_.check("all-server"))
    {
      index_ = spawn<index>(
          archive_,
          (config_.get<fs::path>("directory") / "index").string());
      send(index_, atom("load"));

      LOG(verbose, core) << "publishing index at *:"
          << config_.get<unsigned>("index.port");

      publish(index_, config_.get<unsigned>("index.port"));
    }
    else
    {
      LOG(verbose, core) << "connecting to index at "
          << config_.get<std::string>("index.host") << ":"
          << config_.get<unsigned>("index.port");

      index_ = remote_actor(
          config_.get<std::string>("index.host"),
          config_.get<unsigned>("index.port"));

      LOG(verbose, core) << "connected to index actor @" << index_->id();
    }

    if (config_.check("ingestor-actor"))
    {
      comm::broccoli::init(config_.check("broccoli-messages"),
                           config_.check("broccoli-calltrace"));

      ingestor_ = spawn<ingestor>(tracker_, archive_, index_);
      self->monitor(ingestor_);

      send(ingestor_, atom("initialize"),
          config_.get<size_t>("ingest.max-events-per-chunk"),
          config_.get<size_t>("ingest.max-segment-size") * 1000000,
          config_.get<size_t>("ingest.batch-size"));

      if (config_.check("ingest.events"))
      {
        auto host = config_.get<std::string>("ingest.broccoli-host");
        auto port = config_.get<unsigned>("ingest.broccoli-port");
        auto events = config_.get<std::vector<std::string>>("ingest.broccoli-events");
        send(ingestor_, atom("ingest"), atom("broccoli"), host, port, events);
      }

      if (config_.check("ingest.file-names"))
      {
        auto type = config_.get<std::string>("ingest.file-type");
        auto files = config_.get<std::vector<std::string>>("ingest.file-names");
        for (auto& file : files)
        {
          if (fs::exists(file))
            send(ingestor_, atom("ingest"), type, file);
          else
            LOG(error, core) << "no such file: " << file;
        }
      }

      send(ingestor_, atom("extract"));
    }

    if (config_.check("search-actor") || config_.check("all-server"))
    {
      search_ = spawn<search>(archive_, index_);

      LOG(verbose, core) << "publishing search at *:"
          << config_.get<unsigned>("search.port");

      publish(search_, config_.get<unsigned>("search.port"));
    }
    else if (config_.check("client.expression"))
    {
      LOG(verbose, core) << "connecting to search at "
          << config_.get<std::string>("search.host") << ":"
          << config_.get<unsigned>("search.port");

      search_ = remote_actor(
          config_.get<std::string>("search.host"),
          config_.get<unsigned>("search.port"));

      LOG(verbose, core) << "connected to search actor @" << search_->id();

      auto paginate = config_.get<unsigned>("client.paginate");
      auto& expression = config_.get<std::string>("client.expression");
      query_client_ = spawn<query_client>(search_, expression, paginate);
      self->monitor(query_client_);
      send(query_client_, atom("start"));
    }

    return true;
  }
  catch (network_error const& e)
  {
      LOG(error, core) << "network error: " << e.what();
  }

  return false;
}

void program::stop()
{
  auto shutdown = make_any_tuple(atom("shutdown"));

  if (query_client_)
    query_client_ << shutdown;

  if (config_.check("search-actor") || config_.check("all-server"))
    search_ << shutdown;

  if (config_.check("ingestor-actor"))
    ingestor_ << shutdown;

  if (config_.check("index-actor") || config_.check("all-server"))
    index_ << shutdown;

  if (config_.check("archive-actor") || config_.check("all-server"))
    archive_ << shutdown;

  if (config_.check("tracker-actor") || config_.check("all-server"))
    tracker_ << shutdown;

  // TODO: uncomment once brought back into the game.
  //schema_manager_ << shutdown;

  if (profiler_)
    profiler_ << shutdown;

  if (system_monitor_)
    system_monitor_ << shutdown;

}

} // namespace vast
<|endoftext|>


<commit_before>#include "vast/program.h"

#include <caf/io/all.hpp>
#include <cstdlib>
#include <csignal>
#include <iostream>
#include "vast/archive.h"
#include "vast/exporter.h"
#include "vast/file_system.h"
#include "vast/index.h"
#include "vast/importer.h"
#include "vast/logger.h"
#include "vast/profiler.h"
#include "vast/receiver.h"
#include "vast/serialization.h"
#include "vast/signal_monitor.h"
#include "vast/search.h"
#include "vast/tracker.h"
#include "vast/detail/type_manager.h"
#include "vast/sink/bro.h"
#include "vast/sink/json.h"
#include "vast/source/bro.h"
#include "vast/source/test.h"

#ifdef VAST_HAVE_PCAP
#include "vast/sink/pcap.h"
#include "vast/source/pcap.h"
#endif

#ifdef VAST_HAVE_EDITLINE
#include "vast/console.h"
#endif

using namespace caf;

namespace vast {

program::program(configuration config)
  : config_{std::move(config)}
{
}

message_handler program::make_handler()
{
  attach_functor(
      [=](uint32_t)
      {
        receiver_ = invalid_actor;
        tracker_ = invalid_actor;
        archive_ = invalid_actor;
        index_ = invalid_actor;
        search_ = invalid_actor;
        importer_ = invalid_actor;
        exporter_ = invalid_actor;
      });

  return
  {
    on(atom("run")) >> [=]
    {
      run();
    },
    on(atom("tracker")) >> [=]
    {
      return tracker_;
    },
    on(atom("signal"), arg_match) >> [=](int signal)
    {
      VAST_LOG_ACTOR_VERBOSE("received signal " << signal);
      if (signal == SIGINT || signal == SIGTERM)
        quit(exit::stop);
    },
    on(atom("success")) >> [] { /* nothing to do */ }
  };
}

std::string program::name() const
{
  return "program";
}

void program::run()
{
  auto dir = path{*config_.get("directory")}.complete();

  try
  {
    if (config_.check("core"))
    {
      *config_["receiver"] = true;
      *config_["tracker"] = true;
      *config_["archive"] = true;
      *config_["index"] = true;
      *config_["search"] = true;
    }

    auto monitor = spawn<signal_monitor, detached+linked>(this);
    send(monitor, atom("act"));

    if (config_.check("profiler.rusage")
        || config_.check("profiler.cpu")
        || config_.check("profiler.heap"))
    {
      auto secs = *config_.as<unsigned>("profiler.interval");
      auto prof = spawn<profiler, detached+linked>(
          dir / "log", std::chrono::seconds(secs));

      if (config_.check("profiler.cpu"))
      {
#ifdef VAST_USE_PERFTOOLS_CPU_PROFILER
        send(prof, atom("start"), atom("perftools"), atom("cpu"));
#else
        VAST_LOG_ACTOR_ERROR("not compiled with perftools CPU support");
        quit(exit::error);
        return;
#endif
      }

      if (config_.check("profiler.heap"))
      {
#ifdef VAST_USE_PERFTOOLS_HEAP_PROFILER
        send(prof, atom("start"), atom("perftools"), atom("heap"));
#else
        VAST_LOG_ACTOR_ERROR("not compiled with perftools heap support");
        quit(exit::error);
        return;
#endif
      }

      if (config_.check("profiler.rusage"))
        send(prof, atom("start"), atom("rusage"));
    }

    auto host = *config_.get("tracker.host");
    auto port = *config_.as<unsigned>("tracker.port");
    if (config_.check("tracker"))
    {
      VAST_LOG_ACTOR_INFO("publishes tracker at " <<
                          host << ':' << port);

      tracker_ = spawn<tracker, linked>(dir);
      caf::io::publish(tracker_, port, host.c_str());
    }
    else
    {
      VAST_LOG_ACTOR_INFO("connects to tracker at " <<
                          host << ':' << port);

      tracker_ = caf::io::remote_actor(host, port);
    }

    auto archive_name = *config_.get("archive.name");
    if (config_.check("archive"))
    {
      archive_ = spawn<archive, linked>(
          dir,
          *config_.as<size_t>("archive.max-segments"),
          *config_.as<size_t>("archive.max-segment-size") * 1000000);

      send(tracker_, atom("put"), "archive", archive_, archive_name);
    }

    auto index_name = *config_.get("index.name");
    if (config_.check("index"))
    {
      auto batch_size = *config_.as<size_t>("index.batch-size");
      auto max_events = *config_.as<size_t>("index.max-events");
      auto max_parts = *config_.as<size_t>("index.max-parts");
      auto active_parts = *config_.as<size_t>("index.active-parts");
      index_ = spawn<index, linked>(dir, batch_size, max_events, max_parts,
                                    active_parts);

      send(tracker_, atom("put"), "index", index_, index_name);
    }

    auto receiver_name = *config_.get("receiver.name");
    if (config_.check("receiver"))
    {
      receiver_ = spawn<receiver, linked>();
      send(tracker_, atom("put"), "receiver", receiver_, receiver_name);

      scoped_actor self;
      self->sync_send(tracker_, atom("identifier")).await(
          [=](actor const& identifier)
          {
            send(receiver_, atom("link"), atom("identifier"), identifier);
          });

      if (config_.check("archive"))
      {
        unlink_from(archive_);
        receiver_->link_to(archive_);
        send(tracker_, atom("link"), receiver_name, archive_name);
      }

      if (config_.check("index"))
      {
        unlink_from(index_);
        receiver_->link_to(index_);
        send(tracker_, atom("link"), receiver_name, index_name);
      }
    }

    auto search_name = *config_.get("search.name");
    if (config_.check("search"))
    {
      search_ = spawn<search, linked>();
      send(tracker_, atom("put"), "search", search_, search_name);

      if (config_.check("archive"))
        send(tracker_, atom("link"), search_name, archive_name);

      if (config_.check("index"))
        send(tracker_, atom("link"), search_name, index_name);
    }

    if (auto format = config_.get("importer"))
    {
      auto s = config_.get("import.schema");
      schema sch;
      if (s)
      {
        auto contents = load(*s);
        if (! contents)
        {
          VAST_LOG_ACTOR_ERROR("failed to load schema: " << contents.error());
          quit(exit::error);
          return;
        }

        auto t = to<schema>(*contents);
        if (! t)
        {
          VAST_LOG_ACTOR_ERROR("invalid schema: " << t.error());
          quit(exit::error);
          return;
        }

        sch = *t;
      }

      auto sniff = config_.check("import.sniff-schema");
      auto r = config_.get("import.read");
      actor src;
      if (*format == "pcap")
      {
        if (sniff)
        {
          schema packet_schema;
          packet_schema.add(detail::make_packet_type());
          std::cout << packet_schema << std::flush;
          quit(exit::done);
          return;
        }

#ifdef VAST_HAVE_PCAP
        auto i = config_.get("import.interface");
        auto c = config_.as<size_t>("import.pcap-cutoff");
        auto m = *config_.as<size_t>("import.pcap-maxflows");
        std::string n = i ? *i : *r;
        src = spawn<source::pcap, detached>(sch, std::move(n), c ? *c : -1, m);
#else
        VAST_LOG_ACTOR_ERROR("not compiled with pcap support");
        quit(exit::error);
        return;
#endif
      }
      else if (*format == "bro")
      {
        src = spawn<source::bro, detached>(sch, *r, sniff);
      }
      else if (*format == "test")
      {
        auto id = *config_.as<event_id>("import.test-id");
        auto events = *config_.as<uint64_t>("import.test-events");
        src = spawn<source::test>(sch, id, events);
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("invalid import format: " << *format);
        quit(exit::error);
        return;
      }

      io::compression compression;
      auto method = *config_.get("import.compression");
      if (method == "null")
      {
        compression = io::null;
      }
      else if (method == "lz4")
      {
        compression = io::lz4;
      }
      else if (method == "snappy")
      {
#ifdef VAST_HAVE_SNAPPY
        compression = io::snappy;
#else
        VAST_LOG_ACTOR_ERROR("not compiled with snappy support");
        quit(exit::error);
        return;
#endif
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("unknown compression method");
        quit(exit::error);
        return;
      }

      auto batch_size = *config_.as<uint64_t>("import.batch-size");
      importer_ = spawn<importer, linked>(dir, batch_size, compression);
      send(importer_, atom("source"), src);

      auto importer_name = *config_.get("importer.name");
      send(tracker_, atom("put"), "importer", importer_, importer_name);

      if (config_.check("receiver"))
      {
        // In case we're running in "one-shot" mode where both IMPORTER and
        // RECEIVER share the same program, we initiate the shutdown
        // via IMPORTER to ensure proper delivery of inflight segments from
        // IMPORTER to RECEIVER.
        unlink_from(receiver_);
        importer_->link_to(receiver_);
        send(tracker_, atom("link"), importer_name, receiver_name);
      }
    }
    else if (auto format = config_.get("exporter"))
    {
      auto s = config_.get("export.schema");
      schema sch;
      if (s)
      {
        auto contents = load(*s);
        if (! contents)
        {
          VAST_LOG_ACTOR_ERROR("failed to load schema: " << contents.error());
          quit(exit::error);
          return;
        }

        auto t = to<schema>(*contents);
        if (! t)
        {
          VAST_LOG_ACTOR_ERROR("invalid schema: " << t.error());
          quit(exit::error);
          return;
        }

        sch = *t;
      }

      auto w = config_.get("export.write");
      assert(w);
      actor snk;
      if (*format == "pcap")
      {
#ifdef VAST_HAVE_PCAP
        auto flush = config_.as<uint64_t>("export.pcap-flush");
        assert(flush);
        snk = spawn<sink::pcap, detached>(sch, *w, *flush);
#else
        VAST_LOG_ACTOR_ERROR("not compiled with pcap support");
        quit(exit::error);
        return;
#endif
      }
      else if (*format == "bro")
      {
        snk = spawn<sink::bro>(std::move(*w));
      }
      else if (*format == "json")
      {
        path p{std::move(*w)};
        if (p != "-")
        {
          p = p.complete();
          if (! exists(p.parent()))
          {
            auto t = mkdir(p.parent());
            if (! t)
            {
              VAST_LOG_ACTOR_ERROR("failed to create directory: " <<
                                   p.parent());
              quit(exit::error);
              return;
            }
          }
        }

        snk = spawn<sink::json>(std::move(p));
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("invalid export format: " << *format);
        quit(exit::error);
        return;
      }

      exporter_ = spawn<exporter, linked>();
      send(exporter_, atom("add"), std::move(snk));

      auto exporter_name = *config_.get("exporter.name");
      send(tracker_, atom("put"), "exporter", exporter_, exporter_name);

      auto limit = *config_.as<uint64_t>("export.limit");
      if (limit > 0)
        send(exporter_, atom("limit"), limit);

      auto query = config_.get("export.query");
      assert(query);

      sync_send(tracker_, atom("get"), search_name).then(
          on_arg_match >> [=](error const& e)
          {
            VAST_LOG_ACTOR_ERROR("could not get SEARCH: " << e);
            quit(exit::error);
          },
          [=](actor const& srch)
          {
            sync_send(srch, atom("query"), exporter_, *query).then(
                on_arg_match >> [=](error const& e)
                {
                  VAST_LOG_ACTOR_ERROR("got invalid query: " << e);
                  quit(exit::error);
                },
                [=](expression const& ast, actor qry)
                {
                  VAST_LOG_ACTOR_DEBUG("instantiated query for: " << ast);
                  exporter_->link_to(qry);
                  send(qry, atom("extract"), limit);
                },
                others() >> [=]
                {
                  VAST_LOG_ACTOR_ERROR("got unexpected reply: " <<
                                       to_string(last_dequeued()));

                  quit(exit::error);
                });
          });
    }
    else if (config_.check("console"))
    {
#ifdef VAST_HAVE_EDITLINE
      sync_send(tracker_, atom("get"), search_name).then(
          [=](actor const& search)
          {
            auto c = spawn<console, linked>(search, dir / "console");
            delayed_send(c, std::chrono::milliseconds(200), atom("prompt"));
          },
          [=](error const& e)
          {
            VAST_LOG_ACTOR_ERROR(e);
            quit(exit::error);
          });
#else
      VAST_LOG_ACTOR_ERROR("not compiled with editline support");
      quit(exit::error);
#endif
    }
  }
  catch (network_error const& e)
  {
    VAST_LOG_ACTOR_ERROR("encountered network error: " << e.what());
    quit(exit::error);
  }
}

} // namespace vast
<commit_msg>Fix typo in program config setting.<commit_after>#include "vast/program.h"

#include <caf/io/all.hpp>
#include <cstdlib>
#include <csignal>
#include <iostream>
#include "vast/archive.h"
#include "vast/exporter.h"
#include "vast/file_system.h"
#include "vast/index.h"
#include "vast/importer.h"
#include "vast/logger.h"
#include "vast/profiler.h"
#include "vast/receiver.h"
#include "vast/serialization.h"
#include "vast/signal_monitor.h"
#include "vast/search.h"
#include "vast/tracker.h"
#include "vast/detail/type_manager.h"
#include "vast/sink/bro.h"
#include "vast/sink/json.h"
#include "vast/source/bro.h"
#include "vast/source/test.h"

#ifdef VAST_HAVE_PCAP
#include "vast/sink/pcap.h"
#include "vast/source/pcap.h"
#endif

#ifdef VAST_HAVE_EDITLINE
#include "vast/console.h"
#endif

using namespace caf;

namespace vast {

program::program(configuration config)
  : config_{std::move(config)}
{
}

message_handler program::make_handler()
{
  attach_functor(
      [=](uint32_t)
      {
        receiver_ = invalid_actor;
        tracker_ = invalid_actor;
        archive_ = invalid_actor;
        index_ = invalid_actor;
        search_ = invalid_actor;
        importer_ = invalid_actor;
        exporter_ = invalid_actor;
      });

  return
  {
    on(atom("run")) >> [=]
    {
      run();
    },
    on(atom("tracker")) >> [=]
    {
      return tracker_;
    },
    on(atom("signal"), arg_match) >> [=](int signal)
    {
      VAST_LOG_ACTOR_VERBOSE("received signal " << signal);
      if (signal == SIGINT || signal == SIGTERM)
        quit(exit::stop);
    },
    on(atom("success")) >> [] { /* nothing to do */ }
  };
}

std::string program::name() const
{
  return "program";
}

void program::run()
{
  auto dir = path{*config_.get("directory")}.complete();

  try
  {
    if (config_.check("core"))
    {
      *config_["receiver"] = true;
      *config_["tracker"] = true;
      *config_["archive"] = true;
      *config_["index"] = true;
      *config_["search"] = true;
    }

    auto monitor = spawn<signal_monitor, detached+linked>(this);
    send(monitor, atom("act"));

    if (config_.check("profiler.rusage")
        || config_.check("profiler.cpu")
        || config_.check("profiler.heap"))
    {
      auto secs = *config_.as<unsigned>("profiler.interval");
      auto prof = spawn<profiler, detached+linked>(
          dir / "log", std::chrono::seconds(secs));

      if (config_.check("profiler.cpu"))
      {
#ifdef VAST_USE_PERFTOOLS_CPU_PROFILER
        send(prof, atom("start"), atom("perftools"), atom("cpu"));
#else
        VAST_LOG_ACTOR_ERROR("not compiled with perftools CPU support");
        quit(exit::error);
        return;
#endif
      }

      if (config_.check("profiler.heap"))
      {
#ifdef VAST_USE_PERFTOOLS_HEAP_PROFILER
        send(prof, atom("start"), atom("perftools"), atom("heap"));
#else
        VAST_LOG_ACTOR_ERROR("not compiled with perftools heap support");
        quit(exit::error);
        return;
#endif
      }

      if (config_.check("profiler.rusage"))
        send(prof, atom("start"), atom("rusage"));
    }

    auto host = *config_.get("tracker.host");
    auto port = *config_.as<unsigned>("tracker.port");
    if (config_.check("tracker"))
    {
      VAST_LOG_ACTOR_INFO("publishes tracker at " <<
                          host << ':' << port);

      tracker_ = spawn<tracker, linked>(dir);
      caf::io::publish(tracker_, port, host.c_str());
    }
    else
    {
      VAST_LOG_ACTOR_INFO("connects to tracker at " <<
                          host << ':' << port);

      tracker_ = caf::io::remote_actor(host, port);
    }

    auto archive_name = *config_.get("archive.name");
    if (config_.check("archive"))
    {
      archive_ = spawn<archive, linked>(
          dir,
          *config_.as<size_t>("archive.max-segments"),
          *config_.as<size_t>("archive.max-segment-size") * 1000000);

      send(tracker_, atom("put"), "archive", archive_, archive_name);
    }

    auto index_name = *config_.get("index.name");
    if (config_.check("index"))
    {
      auto batch_size = *config_.as<size_t>("index.batch-size");
      auto max_events = *config_.as<size_t>("index.max-events");
      auto max_parts = *config_.as<size_t>("index.max-parts");
      auto active_parts = *config_.as<size_t>("index.active-parts");
      index_ = spawn<index, linked>(dir, batch_size, max_events, max_parts,
                                    active_parts);

      send(tracker_, atom("put"), "index", index_, index_name);
    }

    auto receiver_name = *config_.get("receiver.name");
    if (config_.check("receiver"))
    {
      receiver_ = spawn<receiver, linked>();
      send(tracker_, atom("put"), "receiver", receiver_, receiver_name);

      scoped_actor self;
      self->sync_send(tracker_, atom("identifier")).await(
          [=](actor const& identifier)
          {
            send(receiver_, atom("link"), atom("identifier"), identifier);
          });

      if (config_.check("archive"))
      {
        unlink_from(archive_);
        receiver_->link_to(archive_);
        send(tracker_, atom("link"), receiver_name, archive_name);
      }

      if (config_.check("index"))
      {
        unlink_from(index_);
        receiver_->link_to(index_);
        send(tracker_, atom("link"), receiver_name, index_name);
      }
    }

    auto search_name = *config_.get("search.name");
    if (config_.check("search"))
    {
      search_ = spawn<search, linked>();
      send(tracker_, atom("put"), "search", search_, search_name);

      if (config_.check("archive"))
        send(tracker_, atom("link"), search_name, archive_name);

      if (config_.check("index"))
        send(tracker_, atom("link"), search_name, index_name);
    }

    if (auto format = config_.get("importer"))
    {
      auto s = config_.get("import.schema");
      schema sch;
      if (s)
      {
        auto contents = load(*s);
        if (! contents)
        {
          VAST_LOG_ACTOR_ERROR("failed to load schema: " << contents.error());
          quit(exit::error);
          return;
        }

        auto t = to<schema>(*contents);
        if (! t)
        {
          VAST_LOG_ACTOR_ERROR("invalid schema: " << t.error());
          quit(exit::error);
          return;
        }

        sch = *t;
      }

      auto sniff = config_.check("import.sniff-schema");
      auto r = config_.get("import.read");
      actor src;
      if (*format == "pcap")
      {
        if (sniff)
        {
          schema packet_schema;
          packet_schema.add(detail::make_packet_type());
          std::cout << packet_schema << std::flush;
          quit(exit::done);
          return;
        }

#ifdef VAST_HAVE_PCAP
        auto i = config_.get("import.interface");
        auto c = config_.as<size_t>("import.pcap-cutoff");
        auto m = *config_.as<size_t>("import.pcap-maxflows");
        std::string n = i ? *i : *r;
        src = spawn<source::pcap, detached>(sch, std::move(n), c ? *c : -1, m);
#else
        VAST_LOG_ACTOR_ERROR("not compiled with pcap support");
        quit(exit::error);
        return;
#endif
      }
      else if (*format == "bro")
      {
        src = spawn<source::bro, detached>(sch, *r, sniff);
      }
      else if (*format == "test")
      {
        auto id = *config_.as<event_id>("import.test-id");
        auto events = *config_.as<uint64_t>("import.test-events");
        src = spawn<source::test>(sch, id, events);
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("invalid import format: " << *format);
        quit(exit::error);
        return;
      }

      io::compression compression;
      auto method = *config_.get("import.compression");
      if (method == "null")
      {
        compression = io::null;
      }
      else if (method == "lz4")
      {
        compression = io::lz4;
      }
      else if (method == "snappy")
      {
#ifdef VAST_HAVE_SNAPPY
        compression = io::snappy;
#else
        VAST_LOG_ACTOR_ERROR("not compiled with snappy support");
        quit(exit::error);
        return;
#endif
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("unknown compression method");
        quit(exit::error);
        return;
      }

      auto batch_size = *config_.as<uint64_t>("import.batch-size");
      importer_ = spawn<importer, linked>(dir, batch_size, compression);
      send(importer_, atom("source"), src);

      auto importer_name = *config_.get("import.name");
      send(tracker_, atom("put"), "importer", importer_, importer_name);

      if (config_.check("receiver"))
      {
        // In case we're running in "one-shot" mode where both IMPORTER and
        // RECEIVER share the same program, we initiate the shutdown
        // via IMPORTER to ensure proper delivery of inflight segments from
        // IMPORTER to RECEIVER.
        unlink_from(receiver_);
        importer_->link_to(receiver_);
        send(tracker_, atom("link"), importer_name, receiver_name);
      }
    }
    else if (auto format = config_.get("exporter"))
    {
      auto s = config_.get("export.schema");
      schema sch;
      if (s)
      {
        auto contents = load(*s);
        if (! contents)
        {
          VAST_LOG_ACTOR_ERROR("failed to load schema: " << contents.error());
          quit(exit::error);
          return;
        }

        auto t = to<schema>(*contents);
        if (! t)
        {
          VAST_LOG_ACTOR_ERROR("invalid schema: " << t.error());
          quit(exit::error);
          return;
        }

        sch = *t;
      }

      auto w = config_.get("export.write");
      assert(w);
      actor snk;
      if (*format == "pcap")
      {
#ifdef VAST_HAVE_PCAP
        auto flush = config_.as<uint64_t>("export.pcap-flush");
        assert(flush);
        snk = spawn<sink::pcap, detached>(sch, *w, *flush);
#else
        VAST_LOG_ACTOR_ERROR("not compiled with pcap support");
        quit(exit::error);
        return;
#endif
      }
      else if (*format == "bro")
      {
        snk = spawn<sink::bro>(std::move(*w));
      }
      else if (*format == "json")
      {
        path p{std::move(*w)};
        if (p != "-")
        {
          p = p.complete();
          if (! exists(p.parent()))
          {
            auto t = mkdir(p.parent());
            if (! t)
            {
              VAST_LOG_ACTOR_ERROR("failed to create directory: " <<
                                   p.parent());
              quit(exit::error);
              return;
            }
          }
        }

        snk = spawn<sink::json>(std::move(p));
      }
      else
      {
        VAST_LOG_ACTOR_ERROR("invalid export format: " << *format);
        quit(exit::error);
        return;
      }

      exporter_ = spawn<exporter, linked>();
      send(exporter_, atom("add"), std::move(snk));

      auto exporter_name = *config_.get("export.name");
      send(tracker_, atom("put"), "exporter", exporter_, exporter_name);

      auto limit = *config_.as<uint64_t>("export.limit");
      if (limit > 0)
        send(exporter_, atom("limit"), limit);

      auto query = config_.get("export.query");
      assert(query);

      sync_send(tracker_, atom("get"), search_name).then(
          on_arg_match >> [=](error const& e)
          {
            VAST_LOG_ACTOR_ERROR("could not get SEARCH: " << e);
            quit(exit::error);
          },
          [=](actor const& srch)
          {
            sync_send(srch, atom("query"), exporter_, *query).then(
                on_arg_match >> [=](error const& e)
                {
                  VAST_LOG_ACTOR_ERROR("got invalid query: " << e);
                  quit(exit::error);
                },
                [=](expression const& ast, actor qry)
                {
                  VAST_LOG_ACTOR_DEBUG("instantiated query for: " << ast);
                  exporter_->link_to(qry);
                  send(qry, atom("extract"), limit);
                },
                others() >> [=]
                {
                  VAST_LOG_ACTOR_ERROR("got unexpected reply: " <<
                                       to_string(last_dequeued()));

                  quit(exit::error);
                });
          });
    }
    else if (config_.check("console"))
    {
#ifdef VAST_HAVE_EDITLINE
      sync_send(tracker_, atom("get"), search_name).then(
          [=](actor const& search)
          {
            auto c = spawn<console, linked>(search, dir / "console");
            delayed_send(c, std::chrono::milliseconds(200), atom("prompt"));
          },
          [=](error const& e)
          {
            VAST_LOG_ACTOR_ERROR(e);
            quit(exit::error);
          });
#else
      VAST_LOG_ACTOR_ERROR("not compiled with editline support");
      quit(exit::error);
#endif
    }
  }
  catch (network_error const& e)
  {
    VAST_LOG_ACTOR_ERROR("encountered network error: " << e.what());
    quit(exit::error);
  }
}

} // namespace vast
<|endoftext|>