plugins/augmentor/augmentor_base.h

/* augmentor_base.h                                                  -*- C++ -*-
   Jeremy Barnes, 3 March 2012
   Copyright (c) 2012 Datacratic.  All rights reserved.

   Base class for bid request augmentors.
*/

#ifndef __rtb__augmentor_base_h__
#define __rtb__augmentor_base_h__

#include "soa/service/zmq.hpp"
#include "soa/types/id.h"
#include "rtbkit/common/auction.h"
#include "rtbkit/common/augmentation.h"
#include "soa/service/service_base.h"
#include "soa/service/zmq_utils.h"
#include "soa/service/socket_per_thread.h"
#include "soa/service/typed_message_channel.h"
#include "jml/arch/futex.h"
#include "jml/utils/ring_buffer.h"
#include "soa/service/zmq_endpoint.h"

#include <boost/make_shared.hpp>
#include <boost/function.hpp>
#include <boost/thread.hpp>


namespace RTBKIT {

/******************************************************************************/
/* AUGMENTATION REQUEST                                                       */
/******************************************************************************/

struct AugmentationRequest
{
    std::string augmentor;                    // Name of the augmentor
    std::string router;                       // Router to respond to
    Id id;                                    // Auction id
    std::shared_ptr<BidRequest> bidRequest;   // Bid request to augment
    std::vector<std::string> agents;          // Agents availble to bid
    double timeAvailableMs;                   // Time to respond
    Date startTime;                           // Start of the latency timer
};


/*****************************************************************************/
/* AUGMENTOR BASE                                                            */
/*****************************************************************************/

struct Augmentor : public ServiceBase, public MessageLoop {

    Augmentor(const std::string & augmentorName,
                  const std::string & serviceName,
                  std::shared_ptr<ServiceProxies> proxies);

    Augmentor(const std::string & augmentorName,
                  const std::string & serviceName,
                  ServiceBase & parent);

    ~Augmentor();

    void init();
    void start();
    void shutdown();

    void configureAndWait();

    typedef boost::function<void (const AugmentationRequest &)> OnRequest;

    OnRequest onRequest;

    void respond(const AugmentationRequest & request,
                 const AugmentationList & response);

private:
    std::string augmentorName; // This can differ from the servicenName!

    ZmqMultipleNamedClientBusProxy toRouters;

    typedef std::pair<AugmentationRequest, AugmentationList> Response;
    TypedMessageSink<Response> responseQueue;

    void handleRouterMessage(const std::string & router,
                             const std::vector<std::string> & message);
};


/*****************************************************************************/
/* MULTI THREADED AUGMENTOR                                                   */
/*****************************************************************************/

struct MultiThreadedAugmentor : public Augmentor {

    MultiThreadedAugmentor(const std::string & augmentorName,
                               const std::string & serviceName,
                               std::shared_ptr<ServiceProxies> proxies);

    MultiThreadedAugmentor(const std::string & augmentorName,
                               const std::string & serviceName,
                               ServiceBase & parent);

    ~MultiThreadedAugmentor();

    void init(int numThreads);
    void shutdown();

protected:

    virtual void doRequestImpl(const AugmentationRequest &) = 0;

private:
    uint64_t numWithInfo;

    ML::RingBufferSWMR<AugmentationRequest> ringBuffer;

    boost::thread_group workers;

    void runWorker();

    boost::thread_group workerThreads;

    int numThreadsCreated;

    volatile bool shutdown_;

    void pushRequest(const AugmentationRequest &);
};


/******************************************************************************/
/* SYNC AUGMENTOR BASE                                                        */
/******************************************************************************/

struct SyncAugmentor : public MultiThreadedAugmentor
{

    SyncAugmentor(const std::string & augmentorName,
                      const std::string& serviceName,
                      std::shared_ptr<ServiceProxies> proxies)
            : MultiThreadedAugmentor(augmentorName, serviceName, proxies)
    {
        doRequest = boost::bind(&SyncAugmentor::onRequest, this, _1);
    }

    SyncAugmentor(const std::string & augmentorName,
                      const std::string & serviceName,
                      ServiceBase& parent)
        : MultiThreadedAugmentor(augmentorName, serviceName, parent)
    {
        doRequest = boost::bind(&SyncAugmentor::onRequest, this, _1);
    }

    boost::function<AugmentationList(const AugmentationRequest &) > doRequest;

    virtual AugmentationList
    onRequest(const AugmentationRequest &)
    {
        throw ML::Exception("onRequest or doRequest must be overridden");
    }

protected:

    virtual void
    doRequestImpl(const AugmentationRequest & request)
    {
        AugmentationList response = doRequest(request);
        respond(request, response);
    }
};


/******************************************************************************/
/* ASYNC AUGMENTOR BASE                                                       */
/******************************************************************************/

struct AsyncAugmentor : public MultiThreadedAugmentor
{

    AsyncAugmentor(const std::string & augmentorName,
                       const std::string & serviceName,
                       std::shared_ptr<ServiceProxies> proxies) :
        MultiThreadedAugmentor(augmentorName, serviceName, proxies)
    {
        doRequest = boost::bind(&AsyncAugmentor::onRequest, this, _1, _2);
    }

    AsyncAugmentor(const std::string & augmentorName,
                       const std::string & serviceName,
                       ServiceBase& parent) :
        MultiThreadedAugmentor(augmentorName, serviceName, parent)
    {
        doRequest = boost::bind(&AsyncAugmentor::onRequest, this, _1, _2);
    }

    typedef std::function<void (const AugmentationList &)> SendResponseCB;

    boost::function<void (const AugmentationRequest &, SendResponseCB) >
        doRequest;

    virtual void
    onRequest(const AugmentationRequest & request, SendResponseCB sendResponse)
    {
        throw ML::Exception("onRequest or doRequest must be overridden");
    };

protected:

    virtual void
    doRequestImpl(const AugmentationRequest & request)
    {
        auto sendResponse = [=](const AugmentationList & response) {
            respond(request, response);
        };
        doRequest(request, sendResponse);
    }
};

} // namespace RTBKIT


#endif /* __rtb__augmentor_base_h__ */