<file_sep>// Copyright (C) 2012  Internet Systems Consortium, Inc. ("ISC")
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS.  IN NO EVENT SHALL ISC BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
// OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.

#include <datasrc/master_loader_callbacks.h>
#include <datasrc/memory/zone_data_loader.h>
#include <datasrc/memory/zone_data_updater.h>
#include <datasrc/memory/logger.h>
#include <datasrc/memory/segment_object_holder.h>
#include <datasrc/memory/util_internal.h>
#include <datasrc/memory/rrset_collection.h>
#include <datasrc/memory/treenode_rrset.h>
#include <datasrc/client.h>

#include <dns/labelsequence.h>
#include <dns/master_loader.h>
#include <dns/rrcollator.h>
#include <dns/rdataclass.h>
#include <dns/rrset.h>
#include <dns/zone_checker.h>

#include <boost/foreach.hpp>
#include <boost/bind.hpp>
#include <boost/optional.hpp>
#include <boost/noncopyable.hpp>

#include <map>

using namespace bundy::dns;
using namespace bundy::dns::rdata;

namespace bundy {
namespace datasrc {
namespace memory {

using detail::SegmentObjectHolder;
using detail::getCoveredType;

namespace { // unnamed namespace

// A helper internal class for \c ZoneDataLoader().  make it non-copyable
// to avoid accidental copy.
//
// The current internal implementation no longer expects that both a
// normal (non RRSIG) RRset and (when signed) its RRSIG are added at
// once, but we do that here anyway to avoid merging RdataSets every
// single time which can be inefficient.
//
// We hold all RRsets of the same owner name in node_rrsets_ and
// node_rrsigsets_, and add/remove the matching pairs of RRsets to the zone
// when we see a new owner name. We do this to limit the size of
// NodeRRsets below. However, RRsets can occur in any order.
//
// The caller is responsible for adding/removing the RRsets of the last group
// in the input sequence by explicitly calling flushNodeRRsets() at the
// end.  It's cleaner and more robust if we let the destructor of this class
// do it, but since we cannot guarantee the adding/removing operation is
// exception free, we don't choose that option to maintain the common
// expectation for destructors.
class ZoneDataLoaderHelper : boost::noncopyable {
public:
    enum OP_MODE {ADD, DELETE};

    // A functor type used for loading.
    typedef boost::function<void(bundy::dns::ConstRRsetPtr, OP_MODE)>
    LoadCallback;

    ZoneDataLoaderHelper(util::MemorySegment& mem_sgmt,
                         const bundy::dns::RRClass& rrclass,
                         const bundy::dns::Name& zone_name,
                         ZoneData& zone_data) :
        updater_(mem_sgmt, rrclass, zone_name, zone_data)
    {}

    void updateFromLoad(const bundy::dns::ConstRRsetPtr& rrset, OP_MODE mode);
    void flushNodeRRsets();

private:
    typedef std::map<bundy::dns::RRType, bundy::dns::ConstRRsetPtr> NodeRRsets;
    typedef NodeRRsets::value_type NodeRRsetsVal;

    // A helper to identify the covered type of an RRSIG.
    const bundy::dns::Name& getCurrentName() const;

private:
    NodeRRsets node_rrsets_;
    NodeRRsets node_rrsigsets_;
    std::vector<bundy::dns::ConstRRsetPtr> non_consecutive_rrsets_;
    ZoneDataUpdater updater_;
    boost::optional<OP_MODE> current_mode_;
};

void
ZoneDataLoaderHelper::updateFromLoad(const ConstRRsetPtr& rrset, OP_MODE mode) {
    // Set current mode.  If the mode is changing, we first need to flush all
    // changes in the previous mode.
    if (current_mode_ && *current_mode_ != mode) {
        flushNodeRRsets();
    }
    current_mode_ = mode;

    // If we see a new name, flush the temporary holders, adding or removing the
    // pairs of RRsets and RRSIGs of the previous name to the zone.
    if ((!node_rrsets_.empty() || !node_rrsigsets_.empty() ||
         !non_consecutive_rrsets_.empty()) &&
        (getCurrentName() != rrset->getName())) {
        flushNodeRRsets();
    }

    // Store this RRset until it can be added/removed in the zone. If an rrtype
    // that's already been seen is found, queue it in a different vector
    // to be merged later.
    const bool is_rrsig = rrset->getType() == RRType::RRSIG();
    NodeRRsets& node_rrsets = is_rrsig ? node_rrsigsets_ : node_rrsets_;
    const RRType& rrtype = is_rrsig ? getCoveredType(rrset) : rrset->getType();
    if (!node_rrsets.insert(NodeRRsetsVal(rrtype, rrset)).second) {
        non_consecutive_rrsets_.insert(non_consecutive_rrsets_.begin(), rrset);
    }

    if (rrset->getRRsig()) {
        updateFromLoad(rrset->getRRsig(), mode);
    }
}

void
ZoneDataLoaderHelper::flushNodeRRsets() {
    // There has been no add or remove operation.  Then flush is no-op too.
    if (!current_mode_) {
        return;
    }

    boost::function<void(const ConstRRsetPtr&, const ConstRRsetPtr&)> op =
        (*current_mode_ == ADD) ?
        boost::bind(&ZoneDataUpdater::add, &updater_, _1, _2) :
        boost::bind(&ZoneDataUpdater::remove, &updater_, _1, _2);

    BOOST_FOREACH(NodeRRsetsVal val, node_rrsets_) {
        // Identify the corresponding RRSIG for the RRset, if any.  If
        // found add/remove both the RRset and its RRSIG at once.
        ConstRRsetPtr sig_rrset;
        NodeRRsets::iterator sig_it = node_rrsigsets_.find(val.first);
        if (sig_it != node_rrsigsets_.end()) {
            sig_rrset = sig_it->second;
            node_rrsigsets_.erase(sig_it);
        }
        op(val.second, sig_rrset);
    }

    // Normally rrsigsets map should be empty at this point, but it's still
    // possible that an RRSIG that doesn't have covered RRset is added/removed;
    // they still remain in the map.  We add/remove them to the zone separately.
    BOOST_FOREACH(NodeRRsetsVal val, node_rrsigsets_) {
        op(ConstRRsetPtr(), val.second);
    }

    // Add/remove any non-consecutive rrsets too.
    BOOST_FOREACH(ConstRRsetPtr rrset, non_consecutive_rrsets_) {
        if (rrset->getType() == RRType::RRSIG()) {
            op(ConstRRsetPtr(), rrset);
        } else {
            op(rrset, ConstRRsetPtr());
        }
    }

    node_rrsets_.clear();
    node_rrsigsets_.clear();
    non_consecutive_rrsets_.clear();
}

const Name&
ZoneDataLoaderHelper::getCurrentName() const {
    if (!node_rrsets_.empty()) {
        return (node_rrsets_.begin()->second->getName());
    }
    assert(!node_rrsigsets_.empty());
    return (node_rrsigsets_.begin()->second->getName());
}

void
logWarning(const dns::Name* zone_name, const dns::RRClass* rrclass,
           const std::string& reason)
{
    LOG_WARN(logger, DATASRC_MEMORY_CHECK_WARNING).arg(*zone_name).
        arg(*rrclass).arg(reason);
}

void
logError(const dns::Name* zone_name, const dns::RRClass* rrclass,
         const std::string& reason)
{
    LOG_ERROR(logger, DATASRC_MEMORY_CHECK_ERROR).arg(*zone_name).arg(*rrclass).
        arg(reason);
}

// A wrapper for dns::MasterLoader used by ZoneDataLoader.  Essentially
// it converts the two callback types.  Note the mostly redundant wrapper of
// boost::bind.  It converts function<void(ConstRRsetPtr)> to
// function<void(RRsetPtr)> (MasterLoader expects the latter).  SunStudio
// doesn't seem to do this conversion if we just pass 'callback'.
void
masterLoaderWrapper(const char* const filename, const Name& origin,
                    const RRClass& zone_class,
                    ZoneDataLoaderHelper::LoadCallback callback)
{
    bool load_ok = false;       // (we don't use it)
    dns::RRCollator collator(boost::bind(callback, _1,
                                         ZoneDataLoaderHelper::ADD));

    try {
        dns::MasterLoader(filename, origin, zone_class,
                          createMasterLoaderCallbacks(origin, zone_class,
                                                      &load_ok),
                          collator.getCallback()).load();
        collator.flush();
    } catch (const dns::MasterLoaderError& e) {
        bundy_throw(ZoneLoaderException, e.what());
    }
}

// The installer called for ZoneDataLoader using a zone iterator
void
generateRRsetFromIterator(ZoneIterator* iterator,
                          ZoneDataLoaderHelper::LoadCallback callback)
{
    ConstRRsetPtr rrset;
    while ((rrset = iterator->getNextRRset()) != NULL) {
        callback(rrset, ZoneDataLoaderHelper::ADD);
    }
}

ConstRRsetPtr
nextDiff(std::vector<ConstRRsetPtr>::const_iterator& it,
         const std::vector<ConstRRsetPtr>::const_iterator& it_end,
         ZoneJournalReaderPtr& jnl_reader)
{
    ConstRRsetPtr next_diff;    // null by default
    if (it != it_end) {
        next_diff = *it;
        ++it;
    } else if (jnl_reader) {
        next_diff = jnl_reader->getNextDiff();
    }
    return (next_diff);
}

// The installer called for ZoneDataLoader using a zone journal reader.
// It performs some minimal sanity checks on the sequence of data, but
// the basic assumption is that any invalid data mean implementation defect
// (not bad user input) and shouldn't happen anyway.
void
applyDiffs(const std::vector<ConstRRsetPtr>* saved_diffs,
           ZoneJournalReaderPtr jnl_reader,
           ZoneDataLoaderHelper::LoadCallback callback)
{
    enum DIFF_MODE {INIT, ADD, DELETE} mode = INIT;
    ConstRRsetPtr rrset;
    std::vector<ConstRRsetPtr>::const_iterator it = saved_diffs->begin();
    const std::vector<ConstRRsetPtr>::const_iterator it_end =
        saved_diffs->end();
    while ((rrset = nextDiff(it, it_end, jnl_reader))) {
        if (rrset->getType() == RRType::SOA()) {
            mode = (mode == INIT || mode == ADD) ? DELETE : ADD;
        } else if (mode == INIT) {
            // diff sequence doesn't begin with SOA. It means broken journal
            // reader implementation.
            bundy_throw(bundy::Unexpected,
                        "broken journal reader: diff not begin with SOA");
        }
        callback(rrset,
                 (mode == ADD) ?
                 ZoneDataLoaderHelper::ADD : ZoneDataLoaderHelper::DELETE);
    }
    if (mode != ADD) {
        // Diff must end in the add mode (there should at least be one
        // add for the final SOA)
        bundy_throw(bundy::Unexpected, "broken journal reader: incomplete");
    }
}

// The following two are helper to extract SOA serial from the zone data
// or the underlying data source.  The returned object is dynamically allocated
// and must be explicitly deleted by the caller (in practice, the returned
// pointer should be immediately encapsulated in a smart pointer).
// Note: ideally we should be able to use boost::optional for such cases,
// but some versions of g++ seem to have a bug and raise a spurious warning
// for the usage we'd take here.  So, for wider portability we take the less
// safer/cleaner approach.
const dns::Serial*
getSerialFromRRset(const dns::AbstractRRset& rrset) {
    if (rrset.getRdataCount() != 1) {
        // This is basically a bug of the data source implementation, but we
        // should at least crash ourselves.
        bundy_throw(Unexpected,
                    "broken SOA RRset is given to zone data loader: " <<
                    rrset.getRdataCount() << " RDATAs (should be 1)");
    }
    dns::RdataIteratorPtr rdit = rrset.getRdataIterator();
    const dns::rdata::generic::SOA& soa =
        dynamic_cast<const dns::rdata::generic::SOA&>(rdit->getCurrent());
    return (new dns::Serial(soa.getSerial()));
}

const dns::Serial*
getSerialFromZoneData(RRClass rrclass, ZoneData* zone_data) {
    if (zone_data) {
        const ZoneNode* origin_node = zone_data->getOriginNode();
        const RdataSet* rdataset = origin_node->getData();
        rdataset = RdataSet::find(rdataset, RRType::SOA());
        if (rdataset) {
            // explicitly instantiate TreeNodeRRset before passing it to
            // getSerialFromRRset(); some versions of g++ don't seem to treat
            // is as a reference of the base class if we pass a temporary.
            const TreeNodeRRset rrset(rrclass, origin_node, rdataset, false);
            return (getSerialFromRRset(rrset));
        }
    }
    return (NULL);
}

void
validateOldData(const Name& origin, ZoneData* old_data) {
    if (!old_data) {  // need validation only when old data are given.
        return;
    }

    uint8_t buf[LabelSequence::MAX_SERIALIZED_LENGTH];
    const LabelSequence old_name =
        old_data->getOriginNode()->getAbsoluteLabels(buf);
    if (!(old_name == LabelSequence(origin))) {
        bundy_throw(BadValue, "zone data loader is given bad old data: "
                    "origin=" << old_name << ", expecting " << origin);
    }
}

} // end of unnamed namespace

class ZoneDataLoader::ZoneDataLoaderImpl {
public:
    ZoneDataLoaderImpl(util::MemorySegment& mem_sgmt,
                       const dns::RRClass& rrclass,
                       const dns::Name& zone_name,
                       const std::string& zone_file,
                       ZoneData* old_data) :
        mem_sgmt_(mem_sgmt), rrclass_(rrclass), zone_name_(zone_name),
        zone_file_(zone_file), datasrc_client_(NULL), old_data_(old_data),
        old_serial_(getSerialFromZoneData(rrclass, old_data))
    {
        validateOldData(zone_name, old_data);
    }

    ZoneDataLoaderImpl(util::MemorySegment& mem_sgmt,
                       const dns::RRClass& rrclass,
                       const dns::Name& zone_name,
                       const DataSourceClient& datasrc_client,
                       ZoneData* old_data) :
        mem_sgmt_(mem_sgmt), rrclass_(rrclass), zone_name_(zone_name),
        datasrc_client_(&datasrc_client), old_data_(old_data),
        old_serial_(getSerialFromZoneData(rrclass, old_data))
    {
        validateOldData(zone_name, old_data);
    }

    LoadResult doLoad();

    ZoneData* commitDiffs(ZoneData* update_data);

private:
    typedef boost::function<void(ZoneDataLoaderHelper::LoadCallback)>
    RRsetInstaller;
    LoadResult doLoadCommon(
        SegmentObjectHolder<ZoneData, RRClass>* data_holder,
        RRsetInstaller installer);
    ZoneJournalReaderPtr getJournalReader(uint32_t begin, uint32_t end) const;
    void saveDiffs();

    util::MemorySegment& mem_sgmt_;
    const dns::RRClass rrclass_;
    const dns::Name zone_name_;
    const std::string zone_file_;
    const DataSourceClient* const datasrc_client_;
    ZoneData* const old_data_;
    boost::scoped_ptr<const dns::Serial> old_serial_;
    ZoneJournalReaderPtr jnl_reader_;

    // To minimize the risk of hitting an exception from the journal reader
    // in commitDiffs(), we save up to MAX_SAVED_DIFFS_ diff RRs in the
    // load() phase.  While it doesn't guarantee a success in commitDiffs()
    // (if it fails we fall back to invalidate the zone and reload the entire
    // zone), this should work for many cases with small updates (like in
    // dynamic updates).
    static const unsigned int MAX_SAVED_DIFFS_ = 100;
    std::vector<ConstRRsetPtr> saved_diffs_;
};

ZoneDataLoader::LoadResult
ZoneDataLoader::ZoneDataLoaderImpl::doLoad() {
    if (datasrc_client_) {
        ZoneIteratorPtr iterator = datasrc_client_->getIterator(zone_name_);
        if (!iterator) {
            // This shouldn't happen for a compliant implementation of
            // DataSourceClient, but we'll protect ourselves from buggy
            // implementations.
            bundy_throw(Unexpected, "getting loader creator for " << zone_name_
                        << "/" << rrclass_ <<
                        " resulted in Null zone iterator");
        }

        const dns::ConstRRsetPtr new_soarrset = iterator->getSOA();
        if (!new_soarrset) {
            // If the source data source doesn't contain SOA, post-load check
            // will fail anyway, so rejecting loading at this point makes sense.
            bundy_throw(ZoneValidationError, "No SOA found for "
                        << zone_name_ << "/" << rrclass_ << "in " <<
                        datasrc_client_->getDataSourceName());
        }
        boost::scoped_ptr<const dns::Serial> new_serial(
            getSerialFromRRset(*new_soarrset));
        if (old_serial_ && (*old_serial_ == *new_serial)) {
            LOG_DEBUG(logger, DBG_TRACE_BASIC, DATASRC_MEMORY_LOAD_SAME_SERIAL).
                arg(zone_name_).arg(rrclass_).arg(old_serial_->getValue()).
                arg(datasrc_client_->getDataSourceName());
            return (LoadResult(old_data_, false));
        } else if (old_serial_ && (*old_serial_ < *new_serial)) {
            jnl_reader_ = getJournalReader(old_serial_->getValue(),
                                           new_serial->getValue());
            if (jnl_reader_) {
                LOG_DEBUG(logger, DBG_TRACE_BASIC,
                          DATASRC_MEMORY_LOAD_USE_JOURNAL).
                    arg(zone_name_).arg(rrclass_).arg(old_serial_->getValue()).
                    arg(new_serial->getValue()).
                    arg(datasrc_client_->getDataSourceName());
                iterator.reset(); // we don't need the iterator any more.
                saveDiffs();
                return (LoadResult(old_data_, false));
            }
        }

        return (doLoadCommon(NULL, boost::bind(generateRRsetFromIterator,
                                               iterator.get(), _1)));
    } else {
        return (doLoadCommon(NULL, boost::bind(masterLoaderWrapper,
                                               zone_file_.c_str(),
                                               zone_name_, rrclass_, _1)));
    }
}

ZoneJournalReaderPtr
ZoneDataLoader::ZoneDataLoaderImpl::getJournalReader(uint32_t begin,
                                                     uint32_t end) const
{
    try {
        const std::pair<ZoneJournalReader::Result, ZoneJournalReaderPtr>
            result = datasrc_client_->getJournalReader(zone_name_, begin, end);
        return (result.second);
    } catch (const bundy::NotImplemented&) {
        // handle this case just like no journal is available for the serials.
    }
    return (ZoneJournalReaderPtr());
}

void
ZoneDataLoader::ZoneDataLoaderImpl::saveDiffs() {
    int count = 0;
    ConstRRsetPtr rrset;
    while (count++ < MAX_SAVED_DIFFS_ && (rrset = jnl_reader_->getNextDiff())) {
        saved_diffs_.push_back(rrset);
    }
    if (!rrset) {
        jnl_reader_.reset();
        if (saved_diffs_.empty()) {
            // In our expected form of diff sequence, it shouldn't be empty,
            // since there should be at least begin and end SOAs.  Eliminating
            // this case at this point makes the later processing easier.
            bundy_throw(ZoneValidationError,
                        "empty diff sequence is provided for load");
        }
    }
}

ZoneData*
ZoneDataLoader::ZoneDataLoaderImpl::commitDiffs(ZoneData* update_data) {
    // we need 'commit' only when we load diffs
    if (!jnl_reader_ && saved_diffs_.empty()) {
        return (update_data);
    }

    // Constructing SegmentObjectHolder result in MemorySegmentGrown.
    // This needs to be handled at the caller as update_data could now be
    // invalid.  But before propagating the exception, we should release the
    // data because the caller has the ownership and we shouldn't destroy it.
    boost::scoped_ptr<SegmentObjectHolder<ZoneData, RRClass> >
        holder(new SegmentObjectHolder<ZoneData, RRClass>(mem_sgmt_, rrclass_));
    try {
        holder->set(update_data);
        // If doLoadCommon returns the holder should have released the data.
        return (doLoadCommon(holder.get(), boost::bind(applyDiffs,
                                                       &saved_diffs_,
                                                       jnl_reader_, _1)).first);
    } catch (...) {
        holder->release();
        throw;
    }
}

ZoneDataLoader::LoadResult
ZoneDataLoader::ZoneDataLoaderImpl::doLoadCommon(
    SegmentObjectHolder<ZoneData, RRClass>* data_holder,
    RRsetInstaller installer)
{
    while (true) { // Try as long as it takes to load and grow the segment
        bool created = false;
        try {
            boost::scoped_ptr<SegmentObjectHolder<ZoneData, RRClass> >
                local_holder;
            if (!data_holder) {
                local_holder.reset(new SegmentObjectHolder<ZoneData, RRClass>
                                   (mem_sgmt_, rrclass_));
                local_holder->set(ZoneData::create(mem_sgmt_, zone_name_));
                data_holder = local_holder.get();
            }

            // Nothing from this point on should throw MemorySegmentGrown.
            // It is handled inside here.
            created = true;

            ZoneDataLoaderHelper loader(mem_sgmt_, rrclass_, zone_name_,
                                        *data_holder->get());
            installer(boost::bind(&ZoneDataLoaderHelper::updateFromLoad,
                                  &loader, _1, _2));
            // Add any last RRsets that were left
            loader.flushNodeRRsets();

            const ZoneNode* origin_node = data_holder->get()->getOriginNode();
            const RdataSet* rdataset = origin_node->getData();
            ZoneData* const loaded_data = data_holder->get();
            // If the zone is and NSEC3-signed, check if it has NSEC3PARAM.
            // If not, it may either just go to NSEC3-unsigned, or there's an
            // operational error in that step, depending on whether there's any
            // NSEC3 RRs in the zone.
            if (loaded_data->isNSEC3Signed() &&
                RdataSet::find(rdataset, RRType::NSEC3PARAM()) == NULL) {
                if (loaded_data->getNSEC3Data()->isEmpty()) {
                    // becoming NSEC3-unsigned.
                    LOG_INFO(logger, DATASRC_MEMORY_MEM_NSEC3_UNSIGNED).
                        arg(zone_name_).arg(rrclass_);
                    NSEC3Data* old_n3data = loaded_data->setNSEC3Data(NULL);
                    NSEC3Data::destroy(mem_sgmt_, old_n3data, rrclass_);
                } else {
                    LOG_WARN(logger, DATASRC_MEMORY_MEM_NO_NSEC3PARAM).
                        arg(zone_name_).arg(rrclass_);
                }
            }

            RRsetCollection collection(*loaded_data, rrclass_);
            const dns::ZoneCheckerCallbacks
                callbacks(boost::bind(&logError, &zone_name_, &rrclass_, _1),
                          boost::bind(&logWarning, &zone_name_, &rrclass_, _1));
            if (!dns::checkZone(zone_name_, rrclass_, collection, callbacks)) {
                bundy_throw(ZoneValidationError,
                            "Errors found when validating zone: "
                            << zone_name_ << "/" << rrclass_);
            }

            // Check loaded serial.  Note that if checkZone() passed, we
            // should have SOA in the ZoneData.
            boost::scoped_ptr<const dns::Serial> new_serial(
                getSerialFromZoneData(rrclass_, loaded_data));
            if (old_serial_ && *old_serial_ >= *new_serial) {
                LOG_WARN(logger, DATASRC_MEMORY_LOADED_SERIAL_NOT_INCREASED).
                    arg(zone_name_).arg(rrclass_).
                    arg(old_serial_->getValue()).arg(new_serial->getValue());
            }
            LOG_DEBUG(logger, DBG_TRACE_BASIC, DATASRC_MEMORY_LOADED).
                arg(zone_name_).arg(rrclass_).arg(new_serial->getValue()).
                arg(loaded_data->isSigned() ? " (DNSSEC signed)" : "");

            return (LoadResult(data_holder->release(), true));
        } catch (const util::MemorySegmentGrown&) {
            assert(!created);
        }
    }
}

ZoneDataLoader::ZoneDataLoader(util::MemorySegment& mem_sgmt,
                               const dns::RRClass& rrclass,
                               const dns::Name& zone_name,
                               const std::string& zone_file,
                               ZoneData* old_data) :
    impl_(NULL)                 // defer until logging to avoid leak
{
    LOG_DEBUG(logger, DBG_TRACE_BASIC, DATASRC_MEMORY_MEM_LOAD_FROM_FILE).
        arg(zone_name).arg(rrclass).arg(zone_file);

    impl_ = new ZoneDataLoaderImpl(mem_sgmt, rrclass, zone_name, zone_file,
                                   old_data);
}

ZoneDataLoader::ZoneDataLoader(util::MemorySegment& mem_sgmt,
                               const dns::RRClass& rrclass,
                               const dns::Name& zone_name,
                               const DataSourceClient& datasrc_client,
                               ZoneData* old_data) :
    impl_(NULL)
{
    LOG_DEBUG(logger, DBG_TRACE_BASIC, DATASRC_MEMORY_MEM_LOAD_FROM_DATASRC).
        arg(zone_name).arg(rrclass).arg(datasrc_client.getDataSourceName());

    impl_ = new ZoneDataLoaderImpl(mem_sgmt, rrclass, zone_name,
                                   datasrc_client, old_data);
}

ZoneDataLoader::~ZoneDataLoader() {
    delete impl_;
}

ZoneDataLoader::LoadResult
ZoneDataLoader::load() {
    return (impl_->doLoad());
}

ZoneData*
ZoneDataLoader::commit(ZoneData* update_data) {
    return (impl_->commitDiffs(update_data));
}

} // namespace memory
} // namespace datasrc
} // namespace bundy
<file_sep># Copyright (C) 2013  Internet Systems Consortium.
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND INTERNET SYSTEMS CONSORTIUM
# DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
# INTERNET SYSTEMS CONSORTIUM BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
# FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import json
from bundy.datasrc import ConfigurableClientList
import bundy.datasrc
from bundy.memmgr.datasrc_info import SegmentInfo

from bundy.log_messages.libmemmgr_messages import *
from bundy.memmgr.logger import logger

class MemorySegmentBuilder:
    """The builder runs in a different thread in the memory manager. It
    waits for commands from the memory manager, and then executes them
    in the given order sequentially.
    """

    # Internal return code for __reset_segment()
    __RESET_SEGMENT_OK = 0
    __RESET_SEGMENT_FAILED = 1
    __RESET_SEGMENT_CREATED = 2

    def __init__(self, sock, cv, command_queue, response_queue):
        """ The constructor takes the following arguments:

            sock: A socket using which this builder object notifies the
                  main thread that it has a response waiting for it.

            cv: A condition variable object that is used by the main
                thread to tell this builder object that new commands are
                available to it. Note that this is also used for
                synchronizing access to the queues, so code that uses
                MemorySegmentBuilder must use this condition variable's
                lock object to synchronize its access to the queues.

            command_queue: A list of commands sent by the main thread to
                           this object. Commands should be executed
                           sequentially in the given order by this
                           object.

            response_queue: A list of responses sent by this object to
                            the main thread. The format of this is
                            currently not strictly defined. Future
                            tickets will be able to define it based on
                            how it's used.
        """

        self._sock = sock
        self._cv = cv
        self._command_queue = command_queue
        self._response_queue = response_queue
        self._shutdown = False

    def __send_response(self, response_msg):
        with self._cv:
            self._response_queue.append(response_msg)

        # Notify (any main thread) on the socket about a
        # response. Otherwise, the main thread may wait in its
        # loop without knowing there was a problem.
        while self._sock.send(b'x') != 1:
            continue

    def __handle_shutdown(self):
        # This method is called when handling the 'shutdown' command. The
        # following tuple is passed:
        #
        # ('shutdown',)
        self._shutdown = True

    def __handle_bad_command(self, bad_command):
        # A bad command was received. Raising an exception is not useful
        # in this case as we are likely running in a different thread
        # from the main thread which would need to be notified. Instead
        # return this in the response queue.
        logger.error(LIBMEMMGR_BUILDER_BAD_COMMAND_ERROR, bad_command)
        self.__send_response(('bad_command',))
        self._shutdown = True

    def __handle_validate(self, args):
        """Handle 'validate' command.

        Command arguments are the same 'load' except the last one:
        'action': callable without any parameter itself, encapsulating
                  any segment-specific validation logic.  It returns
                  a result of the validation.

        This method simply calls the passed action, and returns the result
        back to the memmgr with other command arguments.  This is run in
        the builder thread simply because it may take time.

        """
        _, dsrc_info, rrclass, dsrc_name, action = args

        logger.debug(logger.DBGLVL_TRACE_BASIC,
                     LIBMEMMGR_BUILDER_SEGMENT_VALIDATE, dsrc_name, rrclass)
        try:
            result = action()
        except Exception as ex:
            logger.error(LIBMEMMGR_BUILDER_SEGMENT_VALIDATE_FAIL, dsrc_name,
                         rrclass, ex)
            result = False
        self.__send_response(('validate-completed', dsrc_info, rrclass,
                              dsrc_name, result))

    def __reset_segment(self, clist, dsrc_name, rrclass, params):
        try:
            clist.reset_memory_segment(dsrc_name,
                                       ConfigurableClientList.READ_WRITE,
                                       params)
            logger.debug(logger.DBGLVL_TRACE_BASIC,
                         LIBMEMMGR_BUILDER_SEGMENT_RESET, dsrc_name, rrclass)
            return self.__RESET_SEGMENT_OK
        except Exception as ex:
            logger.error(LIBMEMMGR_BUILDER_RESET_SEGMENT_ERROR, dsrc_name,
                         rrclass, ex)
        try:
            clist.reset_memory_segment(dsrc_name, ConfigurableClientList.CREATE,
                                       params)
            logger.info(LIBMEMMGR_BUILDER_SEGMENT_CREATED, dsrc_name, rrclass)
            return self.__RESET_SEGMENT_CREATED
        except Exception as ex:
            logger.error(LIBMEMMGR_BUILDER_SEGMENT_CREATE_ERROR, dsrc_name,
                         rrclass, ex)
        return self.__RESET_SEGMENT_FAILED

    def _handle_load(self, zone_name, dsrc_info, rrclass, dsrc_name):
        # This method is called when handling the 'load' command. The
        # following tuple is passed:
        #
        # ('load', zone_name, dsrc_info, rrclass, dsrc_name)
        #
        # where:
        #
        #  * zone_name is None or bundy.dns.Name, specifying the zone name
        #    to load. If it's None, it means all zones to be cached in
        #    the specified data source (used for initialization).
        #
        #  * dsrc_info is a DataSrcInfo object corresponding to the
        #    generation ID of the set of data sources for this loading.
        #
        #  * rrclass is an bundy.dns.RRClass object, the RR class of the
        #    data source.
        #
        #  * dsrc_name is a string, specifying a data source name.
        #
        # This is essentially a 'private' method, but allows tests to call it
        # directly; for other purposes shouldn't be called outside of the class.

        clist = dsrc_info.clients_map[rrclass]
        sgmt_info = dsrc_info.segment_info_map[(rrclass, dsrc_name)]
        params = json.dumps(sgmt_info.get_reset_param(SegmentInfo.WRITER))
        result = self.__reset_segment(clist, dsrc_name, rrclass, params)
        if result == self.__RESET_SEGMENT_FAILED:
            self.__send_response(('load-completed', dsrc_info, rrclass,
                                  dsrc_name, False))
            return

        # If we were told to load a single zone but had to create a new
        # segment, we'll need to load all zones, not just this one.
        if result == self.__RESET_SEGMENT_CREATED and zone_name is not None:
            logger.info(LIBMEMMGR_BUILDER_SEGMENT_LOAD_ALL, zone_name, rrclass,
                        dsrc_name)
            zone_name = None
        if zone_name is not None:
            zones = [(None, zone_name)]
        else:
            zones = clist.get_zone_table_accessor(dsrc_name, True)

        errors = 0
        for _, zname in zones:  # note: don't override zone_name here
            # install empty zone initially
            catch_load_error = (zname is None)
            try:
                result, writer = clist.get_cached_zone_writer(zname,
                                                              catch_load_error,
                                                              dsrc_name)
                if result != ConfigurableClientList.CACHE_STATUS_ZONE_SUCCESS:
                    # handle this with other genuine exceptions below
                    raise bundy.datasrc.Error('result=%d' % result)
            except bundy.datasrc.Error as ex:
                logger.error(LIBMEMMGR_BUILDER_GET_ZONE_WRITER_ERROR,
                             zname, dsrc_name, ex)
                errors += 1
                continue

            try:
                error = writer.load()
                if error is not None:
                    logger.error(LIBMEMMGR_BUILDER_ZONE_WRITER_LOAD_1_ERROR,
                                 zname, dsrc_name, error)
                    errors += 1
                    continue
                writer.install()
            except Exception as e:
                logger.error(LIBMEMMGR_BUILDER_ZONE_WRITER_LOAD_2_ERROR,
                             zname, dsrc_name, e)
                errors += 1
                # fall through to cleanup
            writer.cleanup()

        # need to reset the segment so readers can read it (note: memmgr
        # itself doesn't have to keep it open, but there's currently no
        # public API to just clear the segment).  This 'reset' should succeed,
        # so we'll let any exception be propagated.
        clist.reset_memory_segment(dsrc_name,
                                   ConfigurableClientList.READ_ONLY,
                                   params)

        # At this point, we consider the load a failure only if loading a
        # specific zone has failed.
        succeeded = True if (zone_name is None or errors == 0) else False
        self.__send_response(('load-completed', dsrc_info, rrclass, dsrc_name,
                              succeeded))

    def run(self):
        """ This is the method invoked when the builder thread is
            started.  In this thread, be careful when modifying
            variables passed-by-reference in the constructor. If they
            are reassigned, they will not refer to the main thread's
            objects any longer. Any use of command_queue and
            response_queue must be synchronized by acquiring the lock in
            the condition variable. This method must normally terminate
            only when the 'shutdown' command is sent to it.

        """
        while not self._shutdown:
            with self._cv:
                while not self._command_queue:
                    self._cv.wait()
                # Move the queue content to a local queue. Be careful of
                # not making assignments to reference variables.
                local_command_queue = self._command_queue[:]
                del self._command_queue[:]

            # Run commands passed in the command queue sequentially
            # in the given order.  For now, it only supports the
            # "shutdown" command, which just exits the thread.
            for command_tuple in local_command_queue:
                command = command_tuple[0]
                logger.debug(logger.DBGLVL_TRACE_BASIC,
                             LIBMEMMGR_BUILDER_RECEIVED_COMMAND, command)
                if command == 'validate':
                    self.__handle_validate(command_tuple)
                elif command == 'load':
                    # See the comments for _handle_load() for
                    # details of the tuple passed to the "load"
                    # command.
                    _, zone_name, dsrc_info, rrclass, dsrc_name = \
                        command_tuple
                    self._handle_load(zone_name, dsrc_info, rrclass,
                                      dsrc_name)
                elif command == 'shutdown':
                    self.__handle_shutdown()
                    # When the shutdown command is received, we do
                    # not process any further commands.
                    break
                else:
                    self.__handle_bad_command(command)
                    # When a bad command is received, we do not
                    # process any further commands.
                    break
<file_sep># Copyright (C) 2013  Internet Systems Consortium.
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND INTERNET SYSTEMS CONSORTIUM
# DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
# INTERNET SYSTEMS CONSORTIUM BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
# FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import json
import os
import unittest

from bundy.dns import *
import bundy.config
import bundy.datasrc
import bundy.log
from bundy.server_common.datasrc_clients_mgr import DataSrcClientsMgr
from bundy.memmgr.datasrc_info import *

# Defined for easier tests with DataSrcClientsMgr.reconfigure(), which
# only needs get_value() method
class MockConfigData:
    def __init__(self, data):
        self.__data = data

    def get_value(self, identifier):
        return self.__data[identifier], False

class TestSegmentInfo(unittest.TestCase):
    def setUp(self):
        self.__mapped_file_dir = os.environ['TESTDATA_WRITE_PATH']
        self.__sgmt_info = self.__create_sgmtinfo()
        self.__ver_file = self.__mapped_file_dir + \
            '/zone-IN-0-sqlite3-mapped-vers.json'
        self.__mapped_file_base = self.__mapped_file_dir + \
            '/zone-IN-0-sqlite3-mapped.'

    def __create_sgmtinfo(self):
        return SegmentInfo.create('mapped', 0, RRClass.IN,
                                  'sqlite3', {'mapped_file_dir':
                                              self.__mapped_file_dir})

    def tearDown(self):
        if os.path.exists(self.__ver_file):
            os.unlink(self.__ver_file)

    def __check_sgmt_reset_param(self, user_type, expected_ver, sgmt_info=None):
        """Common check on the return value of get_reset_param() for
        MappedSegmentInfo.

        """
        if sgmt_info is None:
            sgmt_info = self.__sgmt_info # by default, we use the common info

        param = sgmt_info.get_reset_param(user_type)
        if expected_ver is None:
            self.assertEqual(None, param['mapped-file'])
        else:
            self.assertEqual(self.__mapped_file_base + str(expected_ver),
                             param['mapped-file'])

    def test_initial_params(self):
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 1)
        self.__check_sgmt_reset_param(SegmentInfo.READER, None)

        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.INIT)
        self.assertEqual(len(self.__sgmt_info.get_readers()), 0)
        self.assertEqual(len(self.__sgmt_info.get_old_readers()), 0)
        self.assertEqual(len(self.__sgmt_info.get_events()), 0)

        raction, waction = self.__sgmt_info.start_validate()
        self.assertEqual(False, raction())
        self.assertEqual(False, waction())

    def test_init_with_verfile(self):
        # Initialize with versions file, storing non-default versions
        vers = {'reader': 1, 'writer': 0}
        with open(self.__ver_file, 'w') as f:
            f.write(json.dumps(vers) + '\n')
        # Create faked mapped files.  For now, they only need to exist.
        with open(self.__mapped_file_base + '0', 'w'): pass
        with open(self.__mapped_file_base + '1', 'w'): pass

        # re-create the segment info, confirm the expected initial versions.
        self.__sgmt_info = self.__create_sgmtinfo()
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 0)
        self.__check_sgmt_reset_param(SegmentInfo.READER, None)

        # since the files exist, both actions should return True
        raction, waction = self.__sgmt_info.start_validate()
        self.assertEqual(True, raction())
        self.assertEqual(True, waction())

        # If either of the mapped files does't exist, Xaction returns False.
        os.unlink(self.__mapped_file_base + '0')
        self.assertEqual(True, raction()) # shouldn't be changed
        self.assertEqual(False, waction())
        os.unlink(self.__mapped_file_base + '1')
        self.assertEqual(False, raction())
        self.assertEqual(False, waction()) # should be the same result

    def __si_to_rvalidate_state(self):
        # Go to a default starting state
        self.__sgmt_info = SegmentInfo.create('mapped', 0, RRClass.IN,
                                              'sqlite3',
                                              {'mapped_file_dir':
                                                   self.__mapped_file_dir})
        self.__sgmt_info.start_validate()
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.R_VALIDATING)

    def __si_to_ready_state(self):
        self.__si_to_copying_state()
        self.__sgmt_info.complete_update(True)
        # purge any existing readers so we can begin from a fresh 'ready' state
        for r in list(self.__sgmt_info.get_readers()):
            self.__sgmt_info.remove_reader(r)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)

    def __si_to_vsynchronizing_state(self):
        self.__si_to_rvalidate_state()
        self.__sgmt_info.add_event((42,))
        self.__sgmt_info.add_reader(3)
        self.assertIsNone(self.__sgmt_info.complete_validate(True))
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.V_SYNCHRONIZING)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})

    def __si_to_updating_state(self):
        self.__si_to_rvalidate_state()
        self.__sgmt_info.add_event(42)
        self.__sgmt_info.add_event((42,))
        self.assertEqual(self.__sgmt_info.complete_validate(True), 42)
        self.assertTupleEqual(self.__sgmt_info.complete_validate(True), (42, ))
        self.__sgmt_info.add_reader(3)
        self.assertSetEqual(self.__sgmt_info.get_readers(), {3})
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)

    def __si_to_synchronizing_state(self):
        self.__si_to_updating_state()
        self.__sgmt_info.complete_update(True)
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.SYNCHRONIZING)

    def __si_to_copying_state(self):
        self.__si_to_synchronizing_state()
        self.__sgmt_info.sync_reader(3)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.COPYING)

    def test_add_event(self):
        self.assertEqual(len(self.__sgmt_info.get_events()), 0)
        self.__sgmt_info.add_event(None)
        self.assertEqual(len(self.__sgmt_info.get_events()), 1)
        self.assertListEqual(self.__sgmt_info.get_events(), [None])

    def test_add_reader(self):
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), set())
        self.__sgmt_info.add_reader(1)
        self.assertSetEqual(self.__sgmt_info.get_readers(), {1})
        self.__sgmt_info.add_reader(3)
        self.assertSetEqual(self.__sgmt_info.get_readers(), {1, 3})
        self.__sgmt_info.add_reader(2)
        self.assertSetEqual(self.__sgmt_info.get_readers(), {1, 2, 3})

        # adding the same existing reader must throw
        self.assertRaises(SegmentInfoError, self.__sgmt_info.add_reader, (1))
        # but the existing readers must be untouched
        self.assertSetEqual(self.__sgmt_info.get_readers(), {1, 3, 2})

        # none of this touches the old readers
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), set())

    def test_start_update(self):
        # in READY state
        # a) there are no events
        self.__si_to_ready_state()
        e = self.__sgmt_info.start_update()
        self.assertIsNone(e)
        # if there are no events, there is nothing to update
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)

        # b) there are events. this is the same as calling
        # self.__si_to_updating_state(), but let's try to be
        # descriptive.
        self.__si_to_ready_state()
        self.__sgmt_info.add_event((42,))
        e = self.__sgmt_info.start_update()
        self.assertTupleEqual(e, (42,))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)

        # in UPDATING state, it should always return None and not
        # change state.
        self.__si_to_updating_state()
        self.assertIsNone(self.__sgmt_info.start_update())
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)

        # in SYNCHRONIZING state, it should always return None
        # and not change state.
        self.__si_to_synchronizing_state()
        self.assertIsNone(self.__sgmt_info.start_update())
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.SYNCHRONIZING)

        # in COPYING state, it should always return None and not
        # change state.
        self.__si_to_copying_state()
        self.assertIsNone(self.__sgmt_info.start_update())
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.COPYING)

    def test_complete_update(self):
        self.__ver_switched = 0        # number of version switches
        def count_verswitch():
            self.__ver_switched += 1

        # in READY state
        self.__si_to_ready_state()
        self.assertRaises(SegmentInfoError,
                          self.__sgmt_info.complete_update, True)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)

        # in UPDATING state this is the same as calling
        # self.__si_to_synchronizing_state(), but let's try to be
        # descriptive.
        #
        # a) with no events
        self.__si_to_updating_state()
        self.__sgmt_info._switch_versions = lambda: count_verswitch()
        e = self.__sgmt_info.complete_update(True)
        self.assertIsNone(e)
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.SYNCHRONIZING)
        self.assertEqual(1, self.__ver_switched)

        # a') with no events, update fails
        self.__si_to_updating_state()
        self.__sgmt_info._switch_versions = lambda: count_verswitch()
        e = self.__sgmt_info.complete_update(False)
        self.assertIsNone(e)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)
        self.assertEqual(1, self.__ver_switched) # no switch

        # b) with events
        self.__si_to_updating_state()
        self.__sgmt_info._switch_versions = lambda: count_verswitch()
        self.__sgmt_info.add_event((81,))
        e = self.__sgmt_info.complete_update(True)
        self.assertIsNone(e) # old_readers is not empty
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.SYNCHRONIZING)
        self.assertEqual(2, self.__ver_switched)

        # b') with events, update fails
        self.__si_to_updating_state()
        self.__sgmt_info._switch_versions = lambda: count_verswitch()
        self.__sgmt_info.add_event((81,))
        e = self.__sgmt_info.complete_update(False)
        self.assertIsNone(e)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)
        self.assertEqual(2, self.__ver_switched) # no switch

        # c) with no readers, complete_update() from UPDATING must go
        # directly to COPYING state
        self.__si_to_ready_state()
        self.__sgmt_info._switch_versions = lambda: count_verswitch()
        self.__sgmt_info.add_event((42,))
        e = self.__sgmt_info.start_update()
        self.assertTupleEqual(e, (42,))
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)
        e = self.__sgmt_info.complete_update(True)
        self.assertEqual(3, self.__ver_switched) # version still switched
        self.assertTupleEqual(e, (42,))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.COPYING)

        # in SYNCHRONIZING state
        self.__si_to_synchronizing_state()
        self.assertRaises(SegmentInfoError,
                          self.__sgmt_info.complete_update, True)
        self.assertEqual(self.__sgmt_info.get_state(),
                         SegmentInfo.SYNCHRONIZING)

        # in COPYING state.  No version switch in this case.
        self.__si_to_copying_state()
        e = self.__sgmt_info.complete_update(True)
        self.assertIsNone(e)
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)
        self.assertEqual(3, self.__ver_switched)

    def __check_complete_validate(self, validated, with_reader):
        info = self.__create_sgmtinfo()
        info.start_validate()

        info.add_event('wopen') # dummy event, but should work
        info.add_event('1st-load') # ditto
        # Initially, the reader segment is not opened
        self.__check_sgmt_reset_param(SegmentInfo.READER, None, info)
        if with_reader:
            info.add_reader('x')

        # complete open.  next state and the returned event will depend on
        # whether we have an existing reader.
        expected_ev = None if with_reader else 'wopen'
        self.assertEqual(info.complete_validate(validated), expected_ev)
        expected_state = SegmentInfo.V_SYNCHRONIZING if with_reader else \
                         SegmentInfo.W_VALIDATING
        self.assertEqual(info.get_state(), expected_state)

        # If it succeeds, reader version will now be available,
        # should be set to the default (0).
        expected_rver = 0 if validated else None
        self.__check_sgmt_reset_param(SegmentInfo.READER, expected_rver, info)

        # completing validate in W_VALIDATING state change the state to UPDATING
        # for the initial load, regardless of whether validate succeeded.
        if expected_state == SegmentInfo.W_VALIDATING:
            self.assertEqual(info.complete_validate(validated), '1st-load')
            self.assertEqual(info.get_state(), SegmentInfo.UPDATING)

        # Calling complte_validate in other state is a bug, result in exception.
        self.assertRaises(SegmentInfoError, info.complete_validate, True)

    def test_complete_validate(self):
        self.__check_complete_validate(True, False)
        self.__check_complete_validate(False, False)
        self.__check_complete_validate(True, True)
        self.__check_complete_validate(False, True)

    # A helper for test_sync_reader(), it mostly works same for both
    # V_SYNCHRONIZING and SYNCHRONIZING states.
    def __check_sync_reader_in_sync(self, in_validate_stage):
        if in_validate_stage:
            init_fn = lambda: self.__si_to_vsynchronizing_state()
            cur_state = SegmentInfo.V_SYNCHRONIZING
            next_state = SegmentInfo.W_VALIDATING
        else:
            init_fn = lambda: self.__si_to_synchronizing_state()
            cur_state = SegmentInfo.SYNCHRONIZING
            next_state = SegmentInfo.COPYING

        # a) ID is not in old readers set. The following call sets up ID 3
        # to be in the old readers set.
        init_fn()
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertRaises(SegmentInfoError, self.__sgmt_info.sync_reader, (1))
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

        # b) ID is in old readers set, but also in readers set.
        init_fn()
        self.__sgmt_info.add_reader(3)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {3})
        self.assertRaises(SegmentInfoError, self.__sgmt_info.sync_reader, (3))
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

        # c) ID is in old readers set, but not in readers set, and
        # old_readers becomes empty.
        init_fn()
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        e = self.__sgmt_info.sync_reader(3)
        self.assertTupleEqual(e, (42,))
        # the ID should be moved from old readers to readers set
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), set())
        self.assertSetEqual(self.__sgmt_info.get_readers(), {3})
        self.assertEqual(self.__sgmt_info.get_state(), next_state)

        # d) ID is in old readers set, but not in readers set, and
        # old_readers doesn't become empty.
        if in_validate_stage:
            self.__si_to_rvalidate_state()
            self.__sgmt_info.add_reader(3)
            self.__sgmt_info.add_event((42,))
        else:
            self.__si_to_updating_state() # this also adds event of '(42,)'
        self.__sgmt_info.add_reader(4)
        if in_validate_stage:
            self.__sgmt_info.complete_validate(True)
        else:
            self.__sgmt_info.complete_update(True)
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3, 4})
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        e = self.__sgmt_info.sync_reader(3)
        self.assertIsNone(e)
        # the ID should be moved from old readers to readers set
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {4})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {3})
        # we should be left in the same state
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

    def test_sync_reader(self):
        # in other states than *SYNCHRONIZING, it's no-op.
        self.__si_to_ready_state()
        self.assertIsNone(self.__sgmt_info.sync_reader(0))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)

        self.__si_to_updating_state()
        self.assertIsNone(self.__sgmt_info.sync_reader(0))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)

        self.__si_to_copying_state()
        self.assertIsNone(self.__sgmt_info.sync_reader(0))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.COPYING)

        # in V_SYNCHRONIZING state:
        self.__check_sync_reader_in_sync(True)

        # in SYNCHRONIZING state:
        self.__check_sync_reader_in_sync(False)

    def __check_remove_reader_in_sync(self, in_validate_stage):
        if in_validate_stage:
            init_fn = lambda: self.__si_to_vsynchronizing_state()
            cur_state = SegmentInfo.V_SYNCHRONIZING
            next_state = SegmentInfo.W_VALIDATING
        else:
            init_fn = lambda: self.__si_to_synchronizing_state()
            cur_state = SegmentInfo.SYNCHRONIZING
            next_state = SegmentInfo.COPYING

        # a) ID is not in old readers set or readers set.
        init_fn()
        self.__sgmt_info.add_reader(4)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {4})
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        self.assertRaises(SegmentInfoError, self.__sgmt_info.remove_reader, (1))
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

        # b) ID is in readers set.
        init_fn()
        self.__sgmt_info.add_reader(4)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {4})
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        e = self.__sgmt_info.remove_reader(4)
        self.assertIsNone(e)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), set())
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        # we only change state if it was removed from old_readers
        # specifically and it became empty.
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

        # c) ID is in old_readers set and it becomes empty.
        init_fn()
        self.__sgmt_info.add_reader(4)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {4})
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        e = self.__sgmt_info.remove_reader(3)
        self.assertTupleEqual(e, (42,))
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), set())
        self.assertSetEqual(self.__sgmt_info.get_readers(), {4})
        self.assertListEqual(self.__sgmt_info.get_events(), [])
        # we only change state if it was removed from old_readers
        # specifically and it became empty.
        self.assertEqual(self.__sgmt_info.get_state(), next_state)

        # d) ID is in old_readers set and it doesn't become empty.
        if in_validate_stage:
            self.__si_to_rvalidate_state()
            self.__sgmt_info.add_reader(3) # see __check_sync_reader_in_sync
            self.__sgmt_info.add_event((42,)) # ditto
        else:
            self.__si_to_updating_state()
        self.__sgmt_info.add_reader(4)
        if in_validate_stage:
            self.__sgmt_info.complete_validate(True)
        else:
            self.__sgmt_info.complete_update(True)
        self.__sgmt_info.add_reader(5)
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {3, 4})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {5})
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        e = self.__sgmt_info.remove_reader(3)
        self.assertIsNone(e)
        self.assertSetEqual(self.__sgmt_info.get_old_readers(), {4})
        self.assertSetEqual(self.__sgmt_info.get_readers(), {5})
        self.assertListEqual(self.__sgmt_info.get_events(), [(42,)])
        # we only change state if it was removed from old_readers
        # specifically and it became empty.
        self.assertEqual(self.__sgmt_info.get_state(), cur_state)

    def test_remove_reader(self):
        # in READY state, it should return None, unless the specified reader
        # doesn't exist.
        self.__si_to_ready_state()
        self.__sgmt_info.add_reader(4)
        self.assertRaises(SegmentInfoError, self.__sgmt_info.remove_reader, (0))
        self.assertIsNone(self.__sgmt_info.remove_reader(4))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.READY)

        # Same for UPDATING state.
        self.__si_to_updating_state()
        self.__sgmt_info.add_reader(4)
        self.assertRaises(SegmentInfoError, self.__sgmt_info.remove_reader, (0))
        self.assertIsNone(self.__sgmt_info.remove_reader(4))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.UPDATING)

        # Same for COPYING state.
        self.__si_to_copying_state()
        self.__sgmt_info.add_reader(4)
        self.assertRaises(SegmentInfoError, self.__sgmt_info.remove_reader, (0))
        self.assertIsNone(self.__sgmt_info.remove_reader(4))
        self.assertEqual(self.__sgmt_info.get_state(), SegmentInfo.COPYING)

        # in V_SYNCHRONIZING state:
        self.__check_remove_reader_in_sync(True)

        # in SYNCHRONIZING state:
        self.__check_remove_reader_in_sync(False)

    def test_switch_versions(self):
        self.__sgmt_info._switch_versions()
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 0)
        self.__check_sgmt_reset_param(SegmentInfo.READER, 1)

        # after switch, it should be persisted in the versions file, so a
        # newly created segment should honor that.
        sgmt_info = self.__create_sgmtinfo()
        sgmt_info.start_validate()
        sgmt_info.complete_validate(True) # to 'validate' the reader version
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 0, sgmt_info)
        self.__check_sgmt_reset_param(SegmentInfo.READER, 1, sgmt_info)

        self.__sgmt_info._switch_versions()
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 1)
        self.__check_sgmt_reset_param(SegmentInfo.READER, 0)

        sgmt_info = self.__create_sgmtinfo()
        sgmt_info.start_validate()
        sgmt_info.complete_validate(True)
        self.__check_sgmt_reset_param(SegmentInfo.WRITER, 1, sgmt_info)
        self.__check_sgmt_reset_param(SegmentInfo.READER, 0, sgmt_info)

    def test_init_others(self):
        # For local type of segment, information isn't needed and won't be
        # created.
        self.assertIsNone(SegmentInfo.create('local', 0, RRClass.IN,
                                             'sqlite3', {}))

        # Unknown type of segment will result in an exception.
        self.assertRaises(SegmentInfoError, SegmentInfo.create, 'unknown', 0,
                          RRClass.IN, 'sqlite3', {})

    def test_missing_methods(self):
        # Bad subclass of SegmentInfo that doesn't implement mandatory methods.
        class TestSegmentInfo(SegmentInfo):
            pass

        self.assertRaises(SegmentInfoError,
                          TestSegmentInfo().get_reset_param,
                          SegmentInfo.WRITER)
        self.assertRaises(SegmentInfoError, TestSegmentInfo()._switch_versions)

    def test_loaded(self):
        self.assertFalse(self.__sgmt_info.loaded())
        self.__si_to_rvalidate_state()
        self.assertFalse(self.__sgmt_info.loaded())
        self.__si_to_vsynchronizing_state()
        self.assertFalse(self.__sgmt_info.loaded())
        self.__si_to_updating_state()
        self.assertFalse(self.__sgmt_info.loaded())

        # It's considered 'loaded' on the first transition to SYNCHRONIZING,
        # and after that it's always so.
        self.__si_to_synchronizing_state()
        self.assertTrue(self.__sgmt_info.loaded())
        self.__si_to_copying_state()
        self.assertTrue(self.__sgmt_info.loaded())

class MockClientList:
    """A mock ConfigurableClientList class.

    Just providing minimal shortcut interfaces needed for DataSrcInfo class.

    """
    def __init__(self, status_list):
        self.__status_list = status_list

    def get_status(self):
        return self.__status_list

class TestDataSrcInfo(unittest.TestCase):
    def setUp(self):
        self.__mapped_file_dir = os.environ['TESTDATA_WRITE_PATH']
        self.__mgr_config = {'mapped_file_dir': self.__mapped_file_dir}
        self.__sqlite3_dbfile = os.environ['TESTDATA_WRITE_PATH'] + '/' + 'zone.db'
        self.__clients_map = {
            # mixture of 'local' and 'mapped' and 'unused' (type =None)
            # segments
            RRClass.IN: MockClientList([('datasrc1', 'local', None),
                                        ('datasrc2', 'mapped', None),
                                        ('datasrc3', None, None)]),
            RRClass.CH: MockClientList([('datasrc2', 'mapped', None),
                                        ('datasrc1', 'local', None)]) }

    def tearDown(self):
        if os.path.exists(self.__sqlite3_dbfile):
            os.unlink(self.__sqlite3_dbfile)

    def __check_sgmt_reset_param(self, sgmt_info, reader_file, writer_file):
        # Check if the initial state of (mapped) segment info object has
        # expected values.
        param = sgmt_info.get_reset_param(SegmentInfo.READER)
        self.assertEqual(reader_file, param['mapped-file'])
        param = sgmt_info.get_reset_param(SegmentInfo.WRITER)
        self.assertEqual(writer_file, param['mapped-file'])

    def test_init(self):
        """Check basic scenarios of constructing DataSrcInfo."""

        # This checks that all data sources of all RR classes are covered,
        # "local" segments are ignored, info objects for "mapped" segments
        # are created and stored in segment_info_map.
        datasrc_info = DataSrcInfo(42, self.__clients_map, self.__mgr_config)
        self.assertEqual(42, datasrc_info.gen_id)
        self.assertEqual(self.__clients_map, datasrc_info.clients_map)
        self.assertEqual(2, len(datasrc_info.segment_info_map))
        sgmt_info = datasrc_info.segment_info_map[(RRClass.IN, 'datasrc2')]

        # by default, reader mapped file isn't available yet, but writer's
        # version is 1.
        base = self.__mapped_file_dir + '/zone-IN-42-datasrc2-mapped.'
        self.__check_sgmt_reset_param(sgmt_info, None, base + '1')
        sgmt_info = datasrc_info.segment_info_map[(RRClass.CH, 'datasrc2')]
        base = self.__mapped_file_dir + '/zone-CH-42-datasrc2-mapped.'
        self.__check_sgmt_reset_param(sgmt_info, None, base + '1')

        # A case where clist.get_status() returns an empty list; shouldn't
        # cause disruption
        self.__clients_map = { RRClass.IN: MockClientList([])}
        datasrc_info = DataSrcInfo(42, self.__clients_map, self.__mgr_config)
        self.assertEqual(42, datasrc_info.gen_id)
        self.assertEqual(0, len(datasrc_info.segment_info_map))

        # A case where clients_map is empty; shouldn't cause disruption
        self.__clients_map = {}
        datasrc_info = DataSrcInfo(42, self.__clients_map, self.__mgr_config)
        self.assertEqual(42, datasrc_info.gen_id)
        self.assertEqual(0, len(datasrc_info.segment_info_map))

    # This test uses real "mmaped" segment and doesn't work without shared
    # memory support.
    @unittest.skipIf(os.environ['HAVE_SHARED_MEMORY'] != 'yes',
                     'shared memory support is not available')
    def test_production(self):
        """Check the behavior closer to a production environment.

        Instead of using mock classes, just for confirming we didn't miss
        something.

        """
        cfg_data = MockConfigData(
            {"_generation_id": 1,
             "classes":
                 {"IN": [{"type": "sqlite3", "cache-enable": True,
                          "cache-type": "mapped", "cache-zones": [],
                          "params": {"database_file": self.__sqlite3_dbfile}}]
                  }
             })
        cmgr = DataSrcClientsMgr(use_cache=True)
        cmgr.reconfigure({}, cfg_data)

        genid, clients_map = cmgr.get_clients_map()
        datasrc_info = DataSrcInfo(genid, clients_map, self.__mgr_config)

        self.assertEqual(1, datasrc_info.gen_id)
        self.assertEqual(clients_map, datasrc_info.clients_map)
        self.assertEqual(1, len(datasrc_info.segment_info_map))
        sgmt_info = datasrc_info.segment_info_map[(RRClass.IN, 'sqlite3')]
        self.assertIsNotNone(sgmt_info.get_reset_param(SegmentInfo.READER))

if __name__ == "__main__":
    bundy.log.init("bundy-test")
    bundy.log.resetUnitTestRootLogger()
    unittest.main()
