examples/augmentor_ex.cc

#include "augmentor_ex.h"

#include "rtbkit/core/agent_configuration/agent_configuration_listener.h"
#include "rtbkit/core/agent_configuration/agent_config.h"
#include "rtbkit/plugins/augmentor/augmentor_base.h"
#include "rtbkit/common/bid_request.h"
#include "soa/service/zmq_named_pub_sub.h"
#include "jml/utils/exc_assert.h"

#include <unordered_map>
#include <mutex>


using namespace std;

namespace RTBKIT {


/******************************************************************************/
/* FREQUENCY CAP STORAGE                                                      */
/******************************************************************************/

struct FrequencyCapStorage
{
    size_t get(const RTBKIT::AccountKey& account, const RTBKIT::UserIds& uids)
    {
        lock_guard<mutex> guard(lock);
        return counters[uids.exchangeId][account[0]];
    }

    void inc(const RTBKIT::AccountKey& account, const RTBKIT::UserIds& uids)
    {
        lock_guard<mutex> guard(lock);
        counters[uids.exchangeId][account[0]]++;
    }

private:

    mutex lock;
    unordered_map<Datacratic::Id, unordered_map<string, size_t> > counters;

};

/******************************************************************************/
/* FREQUENCY CAP AUGMENTOR                                                    */
/******************************************************************************/

FrequencyCapAugmentor::
FrequencyCapAugmentor(
        std::shared_ptr<Datacratic::ServiceProxies> services,
        const string& serviceName,
        const string& augmentorName) :
    SyncAugmentor(augmentorName, serviceName, services),
    storage(new FrequencyCapStorage()),
    agentConfig(getZmqContext()),
    palEvents(getZmqContext())
{
    recordHit("up");
}


void
FrequencyCapAugmentor::
init()
{
    SyncAugmentor::init(2 /* numThreads */);

    /* Manages all the communications with the AgentConfigurationService. */
    agentConfig.init(getServices()->config);
    addSource("FrequencyCapAugmentor::agentConfig", agentConfig);

    /* This lambda will get called when the post auction loop receives a win
       on an auction.
    */
    palEvents.messageHandler = [&] (const vector<zmq::message_t>& msg)
        {
            RTBKIT::AccountKey account(msg[19].toString());
            RTBKIT::UserIds uids =
            RTBKIT::UserIds::createFromJson(msg[15].toString());

            storage->inc(account, uids);
            recordHit("wins");
        };

    palEvents.connectAllServiceProviders(
            "rtbPostAuctionService", "logger", {"MATCHEDWIN"});

    addSource("FrequencyCapAugmentor::palEvents", palEvents);
}


RTBKIT::AugmentationList
FrequencyCapAugmentor::
onRequest(const RTBKIT::AugmentationRequest& request)
{
    recordHit("requests");

    RTBKIT::AugmentationList result;

    const RTBKIT::UserIds& uids = request.bidRequest->userIds;

    for (const string& agent : request.agents) {

        RTBKIT::AgentConfigEntry config = agentConfig.getAgentEntry(agent);

        /* When a new agent comes online there's a race condition where the
           router may send us a bid request for that agent before we receive
           its configuration. This check keeps us safe in that scenario.
        */
        if (!config.valid()) {
            recordHit("unknownConfig");
            continue;
        }

        const RTBKIT::AccountKey& account = config.config->account;

        size_t count = storage->get(account, uids);

        /* The number of times a user has been seen by a given agent can be
           useful to make bid decisions so attach this data to the bid
           request.

           It's also recomended to place your data in an object labeled
           after the augmentor from which it originated.
        */
        result[account[0]].data[request.augmentor] = count;

        /* We tag bid requests that pass the frequency cap filtering because
           if the augmentation doesn't terminate in time or if an error
           occured, then the bid request will not receive any tags and will
           therefor be filtered out for agents that require the frequency
           capping.
        */
        if (count < getCap(request.augmentor, agent, config)) {
            result[account].tags.insert("pass-frequency-cap-ex");
            recordHit("accounts." + account[0] + ".passed");
        }
        else recordHit("accounts." + account[0] + ".capped");
    }

    return result;
}


size_t
FrequencyCapAugmentor::
getCap( const string& augmentor,
        const string& agent,
        const RTBKIT::AgentConfigEntry& config) const
{
    for (const auto& augConfig : config.config->augmentations) {
        if (augConfig.name != augmentor) continue;
        return augConfig.config.asInt();
    }

    /* There's a race condition here where if an agent removes our augmentor
       from its configuration while there are bid requests being augmented
       for that agent then we may not find its config. A sane default is
       good to have in this scenario.
    */

    return 0;
}

} // namespace RTBKIT