<commit_before>#include "globals.hh"
#include "gc.hh"
#include "build.hh"
#include "pathlocks.hh"

#include <boost/shared_ptr.hpp>

#include <sys/types.h>
#include <sys/stat.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>

static string gcLockName = "gc.lock";
static string tempRootsDir = "temproots";
static string gcRootsDir = "gcroots";

/* Acquire the global GC lock.  This is used to prevent new Nix
   processes from starting after the temporary root files have been
   read.  To be precise: when they try to create a new temporary root
   file, they will block until the garbage collector has finished /
   yielded the GC lock. */
static int openGCLock(LockType lockType)
{
    Path fnGCLock = (format("%1%/%2%")
        % nixStateDir % gcLockName).str();

    debug(format("acquiring global GC lock `%1%'") % fnGCLock);

    AutoCloseFD fdGCLock = open(fnGCLock.c_str(), O_RDWR | O_CREAT, 0600);
    if (fdGCLock == -1)
        throw SysError(format("opening global GC lock `%1%'") % fnGCLock);

    lockFile(fdGCLock, lockType, true);

    /* !!! Restrict read permission on the GC root.  Otherwise any
       process that can open the file for reading can DoS the
       collector. */

    return fdGCLock.borrow();
}

void createSymlink(const Path & link, const Path & target, bool careful)
{
    /* Create directories up to `gcRoot'. */
    createDirs(dirOf(link));

    /* Remove the old symlink. */
    if (pathExists(link)) {
        if (careful && (!isLink(link) || !isInStore(readLink(link))))
            throw Error(format("cannot create symlink `%1%'; already exists") % link);
        unlink(link.c_str());
    }

    /* And create the new own. */
    if (symlink(target.c_str(), link.c_str()) == -1)
        throw SysError(format("symlinking `%1%' to `%2%'")
            % link % target);
}

Path addPermRoot(const Path & _storePath, const Path & _gcRoot,
    bool indirect)
{
    Path storePath(canonPath(_storePath));
    Path gcRoot(canonPath(_gcRoot));
    assertStorePath(storePath);

    /* Grab the global GC root.  This prevents the set of permanent
       roots from increasing while a GC is in progress. */
    AutoCloseFD fdGCLock = openGCLock(ltRead);

    if (indirect) {
        string hash = printHash32(hashString(htSHA1, gcRoot));
        Path realRoot = canonPath((format("%1%/%2%/auto/%3%")
            % nixStateDir % gcRootsDir % hash).str());

        {
            SwitchToOriginalUser sw;
            createSymlink(gcRoot, storePath, true);
        }
        createSymlink(realRoot, gcRoot, false);
    }

    else {
        Path rootsDir = canonPath((format("%1%/%2%") % nixStateDir % gcRootsDir).str());

        if (string(gcRoot, 0, rootsDir.size() + 1) != rootsDir + "/")
            throw Error(format(
                "path `%1%' is not a valid garbage collector root; "
                "it's not in the directory `%2%'")
                % gcRoot % rootsDir);

        createSymlink(gcRoot, storePath, false);
    }

    return gcRoot;
}

/* The file to which we write our temporary roots. */
static Path fnTempRoots;
static AutoCloseFD fdTempRoots;

void addTempRoot(const Path & path)
{
    /* Create the temporary roots file for this process. */
    if (fdTempRoots == -1) {

        while (1) {
            Path dir = (format("%1%/%2%") % nixStateDir % tempRootsDir).str();
            createDirs(dir);

            fnTempRoots = (format("%1%/%2%")
                % dir % getpid()).str();

            AutoCloseFD fdGCLock = openGCLock(ltRead);

            fdTempRoots = open(fnTempRoots.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600);
            if (fdTempRoots == -1)
                throw SysError(format("opening temporary roots file `%1%'") % fnTempRoots);

            fdGCLock.close();

            debug(format("acquiring read lock on `%1%'") % fnTempRoots);
            lockFile(fdTempRoots, ltRead, true);

            /* Check whether the garbage collector didn't get in our
               way. */
            struct stat st;
            if (fstat(fdTempRoots, &st) == -1)
                throw SysError(format("statting `%1%'") % fnTempRoots);
            if (st.st_size == 0) break;

            /* The garbage collector deleted this file before we could
               get a lock.  (It won't delete the file after we get a
               lock.)  Try again. */
        }

    }

    /* Upgrade the lock to a write lock.  This will cause us to block
       if the garbage collector is holding our lock. */
    debug(format("acquiring write lock on `%1%'") % fnTempRoots);
    lockFile(fdTempRoots, ltWrite, true);

    string s = path + '\0';
    writeFull(fdTempRoots, (const unsigned char *) s.c_str(), s.size());

    /* Downgrade to a read lock. */
    debug(format("downgrading to read lock on `%1%'") % fnTempRoots);
    lockFile(fdTempRoots, ltRead, true);
}

void removeTempRoots()
{
    if (fdTempRoots != -1) {
        fdTempRoots.close();
        unlink(fnTempRoots.c_str());
    }
}

typedef shared_ptr<AutoCloseFD> FDPtr;
typedef list<FDPtr> FDs;

static void readTempRoots(PathSet & tempRoots, FDs & fds)
{
    /* Read the `temproots' directory for per-process temporary root
       files. */
    Strings tempRootFiles = readDirectory(
        (format("%1%/%2%") % nixStateDir % tempRootsDir).str());

    for (Strings::iterator i = tempRootFiles.begin();
         i != tempRootFiles.end(); ++i)
    {
        Path path = (format("%1%/%2%/%3%") % nixStateDir % tempRootsDir % *i).str();

        debug(format("reading temporary root file `%1%'") % path);

        FDPtr fd(new AutoCloseFD(open(path.c_str(), O_RDWR, 0666)));
        if (*fd == -1) {
            /* It's okay if the file has disappeared. */
            if (errno == ENOENT) continue;
            throw SysError(format("opening temporary roots file `%1%'") % path);
        }

        /* Try to acquire a write lock without blocking.  This can
           only succeed if the owning process has died.  In that case
           we don't care about its temporary roots. */
        if (lockFile(*fd, ltWrite, false)) {
            printMsg(lvlError, format("removing stale temporary roots file `%1%'")
                % path);
            unlink(path.c_str());
            writeFull(*fd, (const unsigned char *) "d", 1);
            continue;
        }

        /* Acquire a read lock.  This will prevent the owning process
           from upgrading to a write lock, therefore it will block in
           addTempRoot(). */
        debug(format("waiting for read lock on `%1%'") % path);
        lockFile(*fd, ltRead, true);

        /* Read the entire file. */
        string contents = readFile(*fd);

        /* Extract the roots. */
        unsigned int pos = 0, end;

        while ((end = contents.find((char) 0, pos)) != string::npos) {
            Path root(contents, pos, end - pos);
            debug(format("got temporary root `%1%'") % root);
            assertStorePath(root);
            tempRoots.insert(root);
            pos = end + 1;
        }

        fds.push_back(fd); /* keep open */
    }
}

static void findRoots(const Path & path, bool recurseSymlinks,
    PathSet & roots)
{
    struct stat st;
    if (lstat(path.c_str(), &st) == -1)
        throw SysError(format("statting `%1%'") % path);

    printMsg(lvlVomit, format("looking at `%1%'") % path);

    if (S_ISDIR(st.st_mode)) {
	Strings names = readDirectory(path);
	for (Strings::iterator i = names.begin(); i != names.end(); ++i)
            findRoots(path + "/" + *i, recurseSymlinks, roots);
    }

    else if (S_ISLNK(st.st_mode)) {
        string target = readLink(path);
        Path target2 = absPath(target, dirOf(path));

        if (isInStore(target2)) {
            debug(format("found root `%1%' in `%2%'")
                % target2 % path);
            roots.insert(toStorePath(target2));
        }

        else if (recurseSymlinks) {
            if (pathExists(target2))
                findRoots(target2, false, roots);
            else {
                printMsg(lvlInfo, format("removing stale link from `%1%' to `%2%'") % path % target2);
                /* Note that we only delete when recursing, i.e., when
                   we are still in the `gcroots' tree.  We never
                   delete stuff outside that tree. */
                unlink(path.c_str());
            }
        }
    }
}

static void dfsVisit(const PathSet & paths, const Path & path,
    PathSet & visited, Paths & sorted)
{
    if (visited.find(path) != visited.end()) return;
    visited.insert(path);

    PathSet references;
    if (isValidPath(path))
        queryReferences(noTxn, path, references);

    for (PathSet::iterator i = references.begin();
         i != references.end(); ++i)
        /* Don't traverse into paths that don't exist.  That can
           happen due to substitutes for non-existent paths. */
        if (*i != path && paths.find(*i) != paths.end())
            dfsVisit(paths, *i, visited, sorted);

    sorted.push_front(path);
}

static Paths topoSort(const PathSet & paths)
{
    Paths sorted;
    PathSet visited;
    for (PathSet::const_iterator i = paths.begin(); i != paths.end(); ++i)
        dfsVisit(paths, *i, visited, sorted);
    return sorted;
}

void collectGarbage(GCAction action, PathSet & result)
{
    result.clear();

    bool gcKeepOutputs =
        queryBoolSetting("gc-keep-outputs", false);
    bool gcKeepDerivations =
        queryBoolSetting("gc-keep-derivations", true);

    /* Acquire the global GC root.  This prevents
       a) New roots from being added.
       b) Processes from creating new temporary root files. */
    AutoCloseFD fdGCLock = openGCLock(ltWrite);

    /* Find the roots.  Since we've grabbed the GC lock, the set of
       permanent roots cannot increase now. */
    Path rootsDir = canonPath((format("%1%/%2%") % nixStateDir % gcRootsDir).str());
    PathSet roots;
    findRoots(rootsDir, true, roots);

    if (action == gcReturnRoots) {
        result = roots;
        return;
    }

    /* Determine the live paths which is just the closure of the
       roots under the `references' relation. */
    PathSet livePaths;
    for (PathSet::const_iterator i = roots.begin(); i != roots.end(); ++i)
        computeFSClosure(canonPath(*i), livePaths);

    if (gcKeepDerivations) {
        for (PathSet::iterator i = livePaths.begin();
             i != livePaths.end(); ++i)
        {
            Path deriver = queryDeriver(noTxn, *i);
            if (deriver != "")
                computeFSClosure(deriver, livePaths);
        }
    }

    if (gcKeepOutputs) {
        /* Hmz, identical to storePathRequisites in nix-store. */
        for (PathSet::iterator i = livePaths.begin();
             i != livePaths.end(); ++i)
            if (isDerivation(*i)) {
                Derivation drv = derivationFromPath(*i);
                for (DerivationOutputs::iterator j = drv.outputs.begin();
                     j != drv.outputs.end(); ++j)
                    if (isValidPath(j->second.path))
                        computeFSClosure(j->second.path, livePaths);
            }
    }

    if (action == gcReturnLive) {
        result = livePaths;
        return;
    }

    /* Read the temporary roots.  This acquires read locks on all
       per-process temporary root files.  So after this point no paths
       can be added to the set of temporary roots. */
    PathSet tempRoots;
    FDs fds;
    readTempRoots(tempRoots, fds);

    /* Close the temporary roots.  Note that we *cannot* do this in
       readTempRoots(), because there we may not have all locks yet,
       meaning that an invalid path can become valid (and thus add to
       the references graph) after we have added it to the closure
       (and computeFSClosure() assumes that the presence of a path
       means that it has already been closed). */
    PathSet tempRootsClosed;
    for (PathSet::iterator i = tempRoots.begin(); i != tempRoots.end(); ++i)
        if (isValidPath(*i))
            computeFSClosure(*i, tempRootsClosed);
        else
            tempRootsClosed.insert(*i);

    /* For testing - see tests/gc-concurrent.sh. */
    if (getenv("NIX_DEBUG_GC_WAIT"))
        sleep(2);

    /* After this point the set of roots or temporary roots cannot
       increase, since we hold locks on everything.  So everything
       that is not currently in in `livePaths' or `tempRootsClosed'
       can be deleted. */

    /* Read the Nix store directory to find all currently existing
       paths. */
    Paths storePaths = readDirectory(nixStore);
    PathSet storePaths2;
    for (Paths::iterator i = storePaths.begin(); i != storePaths.end(); ++i)
        storePaths2.insert(canonPath(nixStore + "/" + *i));

    /* Topologically sort them under the `referers' relation.  That
       is, a < b iff a is in referers(b).  This gives us the order in
       which things can be deleted safely. */
    /* !!! when we have multiple output paths per derivation, this
       will not work anymore because we get cycles. */
    storePaths = topoSort(storePaths2);

    /* Try to delete store paths in the topologically sorted order. */
    for (Paths::iterator i = storePaths.begin(); i != storePaths.end(); ++i) {

        debug(format("considering deletion of `%1%'") % *i);

        if (livePaths.find(*i) != livePaths.end()) {
            debug(format("live path `%1%'") % *i);
            continue;
        }

        if (tempRootsClosed.find(*i) != tempRootsClosed.end()) {
            debug(format("temporary root `%1%'") % *i);
            continue;
        }

        debug(format("dead path `%1%'") % *i);
        result.insert(*i);

        AutoCloseFD fdLock;

        if (action == gcDeleteDead) {

            /* Only delete a lock file if we can acquire a write lock
               on it.  That means that it's either stale, or the
               process that created it hasn't locked it yet.  In the
               latter case the other process will detect that we
               deleted the lock, and retry (see pathlocks.cc). */
            if (i->size() >= 5 && string(*i, i->size() - 5) == ".lock") {

                fdLock = open(i->c_str(), O_RDWR);
                if (fdLock == -1) {
                    if (errno == ENOENT) continue;
                    throw SysError(format("opening lock file `%1%'") % *i);
                }

                if (!lockFile(fdLock, ltWrite, false)) {
                    debug(format("skipping active lock `%1%'") % *i);
                    continue;
                }
            }

            printMsg(lvlInfo, format("deleting `%1%'") % *i);

            /* Okay, it's safe to delete. */
            deleteFromStore(*i);

            if (fdLock != -1)
                /* Write token to stale (deleted) lock file. */
                writeFull(fdLock, (const unsigned char *) "d", 1);
        }
    }
}
<commit_msg>* When finding live paths, the deriver need not be valid.<commit_after>#include "globals.hh"
#include "gc.hh"
#include "build.hh"
#include "pathlocks.hh"

#include <boost/shared_ptr.hpp>

#include <sys/types.h>
#include <sys/stat.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>

static string gcLockName = "gc.lock";
static string tempRootsDir = "temproots";
static string gcRootsDir = "gcroots";

/* Acquire the global GC lock.  This is used to prevent new Nix
   processes from starting after the temporary root files have been
   read.  To be precise: when they try to create a new temporary root
   file, they will block until the garbage collector has finished /
   yielded the GC lock. */
static int openGCLock(LockType lockType)
{
    Path fnGCLock = (format("%1%/%2%")
        % nixStateDir % gcLockName).str();

    debug(format("acquiring global GC lock `%1%'") % fnGCLock);

    AutoCloseFD fdGCLock = open(fnGCLock.c_str(), O_RDWR | O_CREAT, 0600);
    if (fdGCLock == -1)
        throw SysError(format("opening global GC lock `%1%'") % fnGCLock);

    lockFile(fdGCLock, lockType, true);

    /* !!! Restrict read permission on the GC root.  Otherwise any
       process that can open the file for reading can DoS the
       collector. */

    return fdGCLock.borrow();
}

void createSymlink(const Path & link, const Path & target, bool careful)
{
    /* Create directories up to `gcRoot'. */
    createDirs(dirOf(link));

    /* Remove the old symlink. */
    if (pathExists(link)) {
        if (careful && (!isLink(link) || !isInStore(readLink(link))))
            throw Error(format("cannot create symlink `%1%'; already exists") % link);
        unlink(link.c_str());
    }

    /* And create the new own. */
    if (symlink(target.c_str(), link.c_str()) == -1)
        throw SysError(format("symlinking `%1%' to `%2%'")
            % link % target);
}

Path addPermRoot(const Path & _storePath, const Path & _gcRoot,
    bool indirect)
{
    Path storePath(canonPath(_storePath));
    Path gcRoot(canonPath(_gcRoot));
    assertStorePath(storePath);

    /* Grab the global GC root.  This prevents the set of permanent
       roots from increasing while a GC is in progress. */
    AutoCloseFD fdGCLock = openGCLock(ltRead);

    if (indirect) {
        string hash = printHash32(hashString(htSHA1, gcRoot));
        Path realRoot = canonPath((format("%1%/%2%/auto/%3%")
            % nixStateDir % gcRootsDir % hash).str());

        {
            SwitchToOriginalUser sw;
            createSymlink(gcRoot, storePath, true);
        }
        createSymlink(realRoot, gcRoot, false);
    }

    else {
        Path rootsDir = canonPath((format("%1%/%2%") % nixStateDir % gcRootsDir).str());

        if (string(gcRoot, 0, rootsDir.size() + 1) != rootsDir + "/")
            throw Error(format(
                "path `%1%' is not a valid garbage collector root; "
                "it's not in the directory `%2%'")
                % gcRoot % rootsDir);

        createSymlink(gcRoot, storePath, false);
    }

    return gcRoot;
}

/* The file to which we write our temporary roots. */
static Path fnTempRoots;
static AutoCloseFD fdTempRoots;

void addTempRoot(const Path & path)
{
    /* Create the temporary roots file for this process. */
    if (fdTempRoots == -1) {

        while (1) {
            Path dir = (format("%1%/%2%") % nixStateDir % tempRootsDir).str();
            createDirs(dir);

            fnTempRoots = (format("%1%/%2%")
                % dir % getpid()).str();

            AutoCloseFD fdGCLock = openGCLock(ltRead);

            fdTempRoots = open(fnTempRoots.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600);
            if (fdTempRoots == -1)
                throw SysError(format("opening temporary roots file `%1%'") % fnTempRoots);

            fdGCLock.close();

            debug(format("acquiring read lock on `%1%'") % fnTempRoots);
            lockFile(fdTempRoots, ltRead, true);

            /* Check whether the garbage collector didn't get in our
               way. */
            struct stat st;
            if (fstat(fdTempRoots, &st) == -1)
                throw SysError(format("statting `%1%'") % fnTempRoots);
            if (st.st_size == 0) break;

            /* The garbage collector deleted this file before we could
               get a lock.  (It won't delete the file after we get a
               lock.)  Try again. */
        }

    }

    /* Upgrade the lock to a write lock.  This will cause us to block
       if the garbage collector is holding our lock. */
    debug(format("acquiring write lock on `%1%'") % fnTempRoots);
    lockFile(fdTempRoots, ltWrite, true);

    string s = path + '\0';
    writeFull(fdTempRoots, (const unsigned char *) s.c_str(), s.size());

    /* Downgrade to a read lock. */
    debug(format("downgrading to read lock on `%1%'") % fnTempRoots);
    lockFile(fdTempRoots, ltRead, true);
}

void removeTempRoots()
{
    if (fdTempRoots != -1) {
        fdTempRoots.close();
        unlink(fnTempRoots.c_str());
    }
}

typedef shared_ptr<AutoCloseFD> FDPtr;
typedef list<FDPtr> FDs;

static void readTempRoots(PathSet & tempRoots, FDs & fds)
{
    /* Read the `temproots' directory for per-process temporary root
       files. */
    Strings tempRootFiles = readDirectory(
        (format("%1%/%2%") % nixStateDir % tempRootsDir).str());

    for (Strings::iterator i = tempRootFiles.begin();
         i != tempRootFiles.end(); ++i)
    {
        Path path = (format("%1%/%2%/%3%") % nixStateDir % tempRootsDir % *i).str();

        debug(format("reading temporary root file `%1%'") % path);

        FDPtr fd(new AutoCloseFD(open(path.c_str(), O_RDWR, 0666)));
        if (*fd == -1) {
            /* It's okay if the file has disappeared. */
            if (errno == ENOENT) continue;
            throw SysError(format("opening temporary roots file `%1%'") % path);
        }

        /* Try to acquire a write lock without blocking.  This can
           only succeed if the owning process has died.  In that case
           we don't care about its temporary roots. */
        if (lockFile(*fd, ltWrite, false)) {
            printMsg(lvlError, format("removing stale temporary roots file `%1%'")
                % path);
            unlink(path.c_str());
            writeFull(*fd, (const unsigned char *) "d", 1);
            continue;
        }

        /* Acquire a read lock.  This will prevent the owning process
           from upgrading to a write lock, therefore it will block in
           addTempRoot(). */
        debug(format("waiting for read lock on `%1%'") % path);
        lockFile(*fd, ltRead, true);

        /* Read the entire file. */
        string contents = readFile(*fd);

        /* Extract the roots. */
        unsigned int pos = 0, end;

        while ((end = contents.find((char) 0, pos)) != string::npos) {
            Path root(contents, pos, end - pos);
            debug(format("got temporary root `%1%'") % root);
            assertStorePath(root);
            tempRoots.insert(root);
            pos = end + 1;
        }

        fds.push_back(fd); /* keep open */
    }
}

static void findRoots(const Path & path, bool recurseSymlinks,
    PathSet & roots)
{
    struct stat st;
    if (lstat(path.c_str(), &st) == -1)
        throw SysError(format("statting `%1%'") % path);

    printMsg(lvlVomit, format("looking at `%1%'") % path);

    if (S_ISDIR(st.st_mode)) {
	Strings names = readDirectory(path);
	for (Strings::iterator i = names.begin(); i != names.end(); ++i)
            findRoots(path + "/" + *i, recurseSymlinks, roots);
    }

    else if (S_ISLNK(st.st_mode)) {
        string target = readLink(path);
        Path target2 = absPath(target, dirOf(path));

        if (isInStore(target2)) {
            debug(format("found root `%1%' in `%2%'")
                % target2 % path);
            roots.insert(toStorePath(target2));
        }

        else if (recurseSymlinks) {
            if (pathExists(target2))
                findRoots(target2, false, roots);
            else {
                printMsg(lvlInfo, format("removing stale link from `%1%' to `%2%'") % path % target2);
                /* Note that we only delete when recursing, i.e., when
                   we are still in the `gcroots' tree.  We never
                   delete stuff outside that tree. */
                unlink(path.c_str());
            }
        }
    }
}

static void dfsVisit(const PathSet & paths, const Path & path,
    PathSet & visited, Paths & sorted)
{
    if (visited.find(path) != visited.end()) return;
    visited.insert(path);

    PathSet references;
    if (isValidPath(path))
        queryReferences(noTxn, path, references);

    for (PathSet::iterator i = references.begin();
         i != references.end(); ++i)
        /* Don't traverse into paths that don't exist.  That can
           happen due to substitutes for non-existent paths. */
        if (*i != path && paths.find(*i) != paths.end())
            dfsVisit(paths, *i, visited, sorted);

    sorted.push_front(path);
}

static Paths topoSort(const PathSet & paths)
{
    Paths sorted;
    PathSet visited;
    for (PathSet::const_iterator i = paths.begin(); i != paths.end(); ++i)
        dfsVisit(paths, *i, visited, sorted);
    return sorted;
}

void collectGarbage(GCAction action, PathSet & result)
{
    result.clear();

    bool gcKeepOutputs =
        queryBoolSetting("gc-keep-outputs", false);
    bool gcKeepDerivations =
        queryBoolSetting("gc-keep-derivations", true);

    /* Acquire the global GC root.  This prevents
       a) New roots from being added.
       b) Processes from creating new temporary root files. */
    AutoCloseFD fdGCLock = openGCLock(ltWrite);

    /* Find the roots.  Since we've grabbed the GC lock, the set of
       permanent roots cannot increase now. */
    Path rootsDir = canonPath((format("%1%/%2%") % nixStateDir % gcRootsDir).str());
    PathSet roots;
    findRoots(rootsDir, true, roots);

    if (action == gcReturnRoots) {
        result = roots;
        return;
    }

    /* Determine the live paths which is just the closure of the
       roots under the `references' relation. */
    PathSet livePaths;
    for (PathSet::const_iterator i = roots.begin(); i != roots.end(); ++i)
        computeFSClosure(canonPath(*i), livePaths);

    if (gcKeepDerivations) {
        for (PathSet::iterator i = livePaths.begin();
             i != livePaths.end(); ++i)
        {
            /* Note that the deriver need not be valid (e.g., if we
               previously ran the collector with `gcKeepDerivations'
               turned off). */
            Path deriver = queryDeriver(noTxn, *i);
            if (deriver != "" && isValidPath(deriver))
                computeFSClosure(deriver, livePaths);
        }
    }

    if (gcKeepOutputs) {
        /* Hmz, identical to storePathRequisites in nix-store. */
        for (PathSet::iterator i = livePaths.begin();
             i != livePaths.end(); ++i)
            if (isDerivation(*i)) {
                Derivation drv = derivationFromPath(*i);
                for (DerivationOutputs::iterator j = drv.outputs.begin();
                     j != drv.outputs.end(); ++j)
                    if (isValidPath(j->second.path))
                        computeFSClosure(j->second.path, livePaths);
            }
    }

    if (action == gcReturnLive) {
        result = livePaths;
        return;
    }

    /* Read the temporary roots.  This acquires read locks on all
       per-process temporary root files.  So after this point no paths
       can be added to the set of temporary roots. */
    PathSet tempRoots;
    FDs fds;
    readTempRoots(tempRoots, fds);

    /* Close the temporary roots.  Note that we *cannot* do this in
       readTempRoots(), because there we may not have all locks yet,
       meaning that an invalid path can become valid (and thus add to
       the references graph) after we have added it to the closure
       (and computeFSClosure() assumes that the presence of a path
       means that it has already been closed). */
    PathSet tempRootsClosed;
    for (PathSet::iterator i = tempRoots.begin(); i != tempRoots.end(); ++i)
        if (isValidPath(*i))
            computeFSClosure(*i, tempRootsClosed);
        else
            tempRootsClosed.insert(*i);

    /* For testing - see tests/gc-concurrent.sh. */
    if (getenv("NIX_DEBUG_GC_WAIT"))
        sleep(2);

    /* After this point the set of roots or temporary roots cannot
       increase, since we hold locks on everything.  So everything
       that is not currently in in `livePaths' or `tempRootsClosed'
       can be deleted. */

    /* Read the Nix store directory to find all currently existing
       paths. */
    Paths storePaths = readDirectory(nixStore);
    PathSet storePaths2;
    for (Paths::iterator i = storePaths.begin(); i != storePaths.end(); ++i)
        storePaths2.insert(canonPath(nixStore + "/" + *i));

    /* Topologically sort them under the `referers' relation.  That
       is, a < b iff a is in referers(b).  This gives us the order in
       which things can be deleted safely. */
    /* !!! when we have multiple output paths per derivation, this
       will not work anymore because we get cycles. */
    storePaths = topoSort(storePaths2);

    /* Try to delete store paths in the topologically sorted order. */
    for (Paths::iterator i = storePaths.begin(); i != storePaths.end(); ++i) {

        debug(format("considering deletion of `%1%'") % *i);

        if (livePaths.find(*i) != livePaths.end()) {
            debug(format("live path `%1%'") % *i);
            continue;
        }

        if (tempRootsClosed.find(*i) != tempRootsClosed.end()) {
            debug(format("temporary root `%1%'") % *i);
            continue;
        }

        debug(format("dead path `%1%'") % *i);
        result.insert(*i);

        AutoCloseFD fdLock;

        if (action == gcDeleteDead) {

            /* Only delete a lock file if we can acquire a write lock
               on it.  That means that it's either stale, or the
               process that created it hasn't locked it yet.  In the
               latter case the other process will detect that we
               deleted the lock, and retry (see pathlocks.cc). */
            if (i->size() >= 5 && string(*i, i->size() - 5) == ".lock") {

                fdLock = open(i->c_str(), O_RDWR);
                if (fdLock == -1) {
                    if (errno == ENOENT) continue;
                    throw SysError(format("opening lock file `%1%'") % *i);
                }

                if (!lockFile(fdLock, ltWrite, false)) {
                    debug(format("skipping active lock `%1%'") % *i);
                    continue;
                }
            }

            printMsg(lvlInfo, format("deleting `%1%'") % *i);

            /* Okay, it's safe to delete. */
            deleteFromStore(*i);

            if (fdLock != -1)
                /* Write token to stale (deleted) lock file. */
                writeFull(fdLock, (const unsigned char *) "d", 1);
        }
    }
}
<|endoftext|>


<commit_before>/*************************************************************************
 *
 *  $RCSfile: wordcountdialog.hxx,v $
 *
 *  $Revision: 1.2 $
 *
 *  last change: $Author: obo $ $Date: 2004-04-27 13:44:33 $
 *
 *  The Contents of this file are made available subject to the terms of
 *  either of the following licenses
 *
 *         - GNU Lesser General Public License Version 2.1
 *         - Sun Industry Standards Source License Version 1.1
 *
 *  Sun Microsystems Inc., October, 2000
 *
 *  GNU Lesser General Public License Version 2.1
 *  =============================================
 *  Copyright 2000 by Sun Microsystems, Inc.
 *  901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *  This library is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU Lesser General Public
 *  License version 2.1, as published by the Free Software Foundation.
 *
 *  This library 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
 *  Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  License along with this library; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *  MA  02111-1307  USA
 *
 *
 *  Sun Industry Standards Source License Version 1.1
 *  =================================================
 *  The contents of this file are subject to the Sun Industry Standards
 *  Source License Version 1.1 (the "License"); You may not use this file
 *  except in compliance with the License. You may obtain a copy of the
 *  License at http://www.openoffice.org/license.html.
 *
 *  Software provided under this License is provided on an "AS IS" basis,
 *  WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING,
 *  WITHOUT LIMITATION, WARRANTIES THAT THE SOFTWARE IS FREE OF DEFECTS,
 *  MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE, OR NON-INFRINGING.
 *  See the License for the specific provisions governing your rights and
 *  obligations concerning the Software.
 *
 *  The Initial Developer of the Original Code is: Sun Microsystems, Inc.
 *
 *  Copyright: 2000 by Sun Microsystems, Inc.
 *
 *  All Rights Reserved.
 *
 *  Contributor(s): _______________________________________
 *
 *
 ************************************************************************/
#ifndef SW_WORDCOUNTDIALOG_HXX
#define SW_WORDCOUNTDIALOG_HXX
#ifndef _BASEDLGS_HXX
#include <sfx2/basedlgs.hxx>
#endif
#ifndef _STDCTRL_HXX
#include <svtools/stdctrl.hxx>
#endif
#ifndef _SV_BUTTON_HXX
#include <vcl/button.hxx>
#endif
struct SwDocStat;
class SwWordCountDialog : public SfxModalDialog
{
    FixedLine       aCurrentFL;
    FixedText       aCurrentWordFT;
    FixedInfo       aCurrentWordFI;
    FixedText       aCurrentCharacterFT;
    FixedInfo       aCurrentCharacterFI;

    FixedLine       aDocFL;
    FixedText       aDocWordFT;
    FixedInfo       aDocWordFI;
    FixedText       aDocCharacterFT;
    FixedInfo       aDocCharacterFI;

    FixedLine       aBottomFL;

    OKButton        aOK;
    HelpButton      aHelp;

public:
    SwWordCountDialog(Window* pParent);
    ~SwWordCountDialog();

    void    SetValues(const SwDocStat& rCurrent, const SwDocStat& rDoc);
};

#endif
<commit_msg>INTEGRATION: CWS ooo19126 (1.2.742); FILE MERGED 2005/09/05 13:45:56 rt 1.2.742.1: #i54170# Change license header: remove SISSL<commit_after>/*************************************************************************
 *
 *  OpenOffice.org - a multi-platform office productivity suite
 *
 *  $RCSfile: wordcountdialog.hxx,v $
 *
 *  $Revision: 1.3 $
 *
 *  last change: $Author: rt $ $Date: 2005-09-09 10:18:42 $
 *
 *  The Contents of this file are made available subject to
 *  the terms of GNU Lesser General Public License Version 2.1.
 *
 *
 *    GNU Lesser General Public License Version 2.1
 *    =============================================
 *    Copyright 2005 by Sun Microsystems, Inc.
 *    901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License version 2.1, as published by the Free Software Foundation.
 *
 *    This library 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
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with this library; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *    MA  02111-1307  USA
 *
 ************************************************************************/
#ifndef SW_WORDCOUNTDIALOG_HXX
#define SW_WORDCOUNTDIALOG_HXX
#ifndef _BASEDLGS_HXX
#include <sfx2/basedlgs.hxx>
#endif
#ifndef _STDCTRL_HXX
#include <svtools/stdctrl.hxx>
#endif
#ifndef _SV_BUTTON_HXX
#include <vcl/button.hxx>
#endif
struct SwDocStat;
class SwWordCountDialog : public SfxModalDialog
{
    FixedLine       aCurrentFL;
    FixedText       aCurrentWordFT;
    FixedInfo       aCurrentWordFI;
    FixedText       aCurrentCharacterFT;
    FixedInfo       aCurrentCharacterFI;

    FixedLine       aDocFL;
    FixedText       aDocWordFT;
    FixedInfo       aDocWordFI;
    FixedText       aDocCharacterFT;
    FixedInfo       aDocCharacterFI;

    FixedLine       aBottomFL;

    OKButton        aOK;
    HelpButton      aHelp;

public:
    SwWordCountDialog(Window* pParent);
    ~SwWordCountDialog();

    void    SetValues(const SwDocStat& rCurrent, const SwDocStat& rDoc);
};

#endif
<|endoftext|>


<commit_before>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifdef INTEL_MKL
#include <string>
#include <unordered_map>
#include "tensorflow/core/common_runtime/eager/eager_op_rewrite_registry.h"
#include "tensorflow/core/graph/mkl_graph_util.h"
#include "tensorflow/core/graph/mkl_layout_pass.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/util/mkl_util.h"
#include "tensorflow/core/util/util.h"

namespace tensorflow {

class MklEagerOpRewrite : public EagerOpRewrite {
 public:
  MklEagerOpRewrite(string name, string file, string line);
  struct MklEagerOp {
    string op_name;
    std::function<bool(EagerOperation*)> RewriteRule;
    std::function<Status(EagerOperation*, std::unique_ptr<EagerOperation>*)>
        CreateMklOp;
  };

 private:
  std::unordered_map<std::string, MklEagerOp> mkl_eager_ops_;

  // The entry point to execute the op rewrite.
  Status Run(EagerOperation* orig_op,
             std::unique_ptr<tensorflow::EagerOperation>* out_op);

  // Initializes the new op and sets up its inputs and attributes
  static Status SetupNewOp(EagerOperation* orig_op, const string mkl_op_name,
                           std::unique_ptr<EagerOperation>* new_mkl_op);

  // Generic rewrite that can be used for any mkl op that doesn't need
  // special processing.
  static Status CreateGenericMklOp(EagerOperation* orig_op,
                                   std::unique_ptr<EagerOperation>* mkl_op);

  // Creates new MKL op for Conv2D, Conv2DBackpropInput and
  // Conv2DBackpropFilter.
  static Status CreateMklConv2DOp(
      EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_conv2d_op);

  // Rewrite rule for Conv2D, Conv2DBackpropInput and Conv2DBackpropFilter.
  static bool RewriteConv2D(EagerOperation* op);

  // Calls op-specific rewrite function to create new MKL op.
  Status RewriteToMklOp(EagerOperation* orig_op,
                        std::unique_ptr<EagerOperation>* mkl_op);

  // Check whether we can rewrite the op to MKL one or not.
  bool ShouldRewriteOp(EagerOperation* op);

  // Default rewrite rule to be used when rewrite should happen without any
  // restriction.
  static bool AlwaysRewrite(EagerOperation* op) { return true; }

  // Check if kernel is registered for a particular op.
  bool FastCheckIfKernelRegistered(std::string op_name, DataType dt);

  // This is called by FastCheckIfKernelRegistered once per unique op name
  // and data type.
  bool SlowCheckIfKernelRegistered(std::string op_name, DataType dt);

  // Helper function to insert mkl_eager_ops to Map
  void InsertMKLEagerOps(MklEagerOp op);

  // Map used by FastCheckIfKernelRegistered.
  std::unordered_map<std::string, bool> registered_kernels_map_;
};

REGISTER_REWRITE(EagerOpRewriteRegistry::PRE_EXECUTION, MklEagerOpRewrite);

// Constructor
MklEagerOpRewrite::MklEagerOpRewrite(string name, string file, string line)
    : EagerOpRewrite(name, file, line) {
  registered_kernels_map_ = std::unordered_map<std::string, bool>();
  InsertMKLEagerOps({"BatchMatMul", AlwaysRewrite, CreateGenericMklOp});
  InsertMKLEagerOps({"BatchMatMulV2", AlwaysRewrite, CreateGenericMklOp});
  InsertMKLEagerOps({"Conv2D", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"Conv2DBackpropInput", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"Conv2DBackpropFilter", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"MatMul", AlwaysRewrite, CreateGenericMklOp});
};

void MklEagerOpRewrite::InsertMKLEagerOps(MklEagerOp op) {
  mkl_eager_ops_.insert(std::make_pair(op.op_name, op));
}

Status MklEagerOpRewrite::Run(
    EagerOperation* orig_op,
    std::unique_ptr<tensorflow::EagerOperation>* out_op) {
  if (ShouldRewriteOp(orig_op)) {
    TF_CHECK_OK(RewriteToMklOp(orig_op, out_op));
  }
  return Status::OK();
}

Status MklEagerOpRewrite::SetupNewOp(
    EagerOperation* orig_op, const string mkl_op_name,
    std::unique_ptr<EagerOperation>* new_mkl_op) {
  const tensorflow::AttrTypeMap* types;
  bool is_function = false;
  TF_RETURN_IF_ERROR(
      tensorflow::AttrTypeMapForOp(mkl_op_name.c_str(), &types, &is_function));
  EagerContext* ctx = orig_op->EagerContext();
  new_mkl_op->reset(new tensorflow::EagerOperation(ctx, mkl_op_name.c_str(),
                                                   is_function, types));

  int num_inputs = orig_op->Inputs().size();
  // Add all inputs to the new op.
  for (int i = 0; i < num_inputs; ++i) {
    (*new_mkl_op)->AddInput(orig_op->Inputs()[i]);
  }

  // Copy all attributes to the new op.
  const NodeDef& orig_ndef = orig_op->MutableAttrs()->BuildNodeDef();

  AttrSlice attr_list(orig_ndef);
  for (const auto& attr : attr_list) {
    (*new_mkl_op)->MutableAttrs()->Set(attr.first, attr.second);
  }

  (*new_mkl_op)
      ->MutableAttrs()
      ->Set("_kernel", mkl_op_registry::kMklNameChangeOpLabel);

  if (orig_op->Device() != nullptr) {
    (*new_mkl_op)->SetDevice(orig_op->Device());
  } else {
    string device_name = orig_op->GetDeviceName();
    (*new_mkl_op)->SetDeviceName(device_name.c_str());
  }
  return Status::OK();
}

Status MklEagerOpRewrite::CreateGenericMklOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_op) {
  const string mkl_op_name = mkl_op_registry::GetMklOpName(orig_op->Name());
  TF_CHECK_OK(SetupNewOp(orig_op, mkl_op_name, mkl_op));
  return Status::OK();
}

Status MklEagerOpRewrite::CreateMklConv2DOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_conv2d_op) {
  const string mkl_op_name =
      mkl_op_registry::GetMklEagerOpName(orig_op->Name());
  TF_CHECK_OK(SetupNewOp(orig_op, mkl_op_name, mkl_conv2d_op));
  return Status::OK();
}

bool MklEagerOpRewrite::ShouldRewriteOp(EagerOperation* op) {
  // Don't rewrite the op if MKL use is disabled at runtime.
  if (DisableMKL()) {
    return false;
  }
  DataType data_type;
  if (op->Attrs().Get("T", &data_type) != Status::OK()) {
    return false;
  }
  // Check if we have registered MKL kernel for this op.
  bool kernel_found = FastCheckIfKernelRegistered(op->Name(), data_type);
  if (!kernel_found) {
    return false;
  }

  // Find and call the op's rewrite rule that determines whether we need to
  // rewrite this op or not.
  auto it = mkl_eager_ops_.find(op->Name());
  if (it != mkl_eager_ops_.end()) {
    // Eager op found so verify Rewrite
    if (it->second.RewriteRule(op)) {
      return true;
    }
  }
  return false;
}

bool MklEagerOpRewrite::FastCheckIfKernelRegistered(std::string op_name,
                                                    DataType dt) {
  // Check for kernel registration only once per op name and data type
  // for performance reasons.
  string registered_kernels_key = op_name + std::to_string(dt);
  auto kernel_element = registered_kernels_map_.find(registered_kernels_key);
  bool kernel_registered = false;
  if (kernel_element == registered_kernels_map_.end()) {
    // Kernel registration is not verified even once yet.
    // So verify and store registration.
    kernel_registered = SlowCheckIfKernelRegistered(op_name, dt);
    registered_kernels_map_.insert(
        std::make_pair(registered_kernels_key, kernel_registered));
  } else {
    // Kernel is visited atleast once. return stored registration result.
    kernel_registered = kernel_element->second;
  }

  return kernel_registered;
}

bool MklEagerOpRewrite::SlowCheckIfKernelRegistered(string op_name,
                                                    DataType dt) {
  // Find if the eager op_name exists in mkl_eager_ops_ list.
  auto element = mkl_eager_ops_.find(op_name);
  if (element != mkl_eager_ops_.end() && dt == DT_FLOAT) {
    // Eager Op exists. So verify registry and return registered or not.
    return (mkl_op_registry::IsMklNameChangeOp(
                mkl_op_registry::GetMklEagerOpName(op_name), dt) ||
            mkl_op_registry::IsMklNameChangeOp(
                mkl_op_registry::GetMklOpName(op_name), dt));
  } else {
    return false;
  }
}

Status MklEagerOpRewrite::RewriteToMklOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_op) {
  // TODO(intel-tf): mkl_eager_ops_ lookup can be reduced from twice
  // (once each in ShouldRewriteOp & RewriteToMklOp) to just once.
  mkl_eager_ops_[orig_op->Name()].CreateMklOp(orig_op, mkl_op);
  return Status::OK();
}

bool MklEagerOpRewrite::RewriteConv2D(EagerOperation* op) {
  const NodeDef& ndef = op->MutableAttrs()->BuildNodeDef();
  string padding;
  TF_CHECK_OK(GetNodeAttr(ndef, "padding", &padding));
  // Right now MKL Conv2D does not support explicit padding.
  return (padding != "EXPLICIT");
}

}  // namespace tensorflow
#endif  // INTEL_MKL
<commit_msg>Initialize variable in initializer list<commit_after>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifdef INTEL_MKL
#include <string>
#include <unordered_map>
#include "tensorflow/core/common_runtime/eager/eager_op_rewrite_registry.h"
#include "tensorflow/core/graph/mkl_graph_util.h"
#include "tensorflow/core/graph/mkl_layout_pass.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/util/mkl_util.h"
#include "tensorflow/core/util/util.h"

namespace tensorflow {

class MklEagerOpRewrite : public EagerOpRewrite {
 public:
  MklEagerOpRewrite(string name, string file, string line);
  struct MklEagerOp {
    string op_name;
    std::function<bool(EagerOperation*)> RewriteRule;
    std::function<Status(EagerOperation*, std::unique_ptr<EagerOperation>*)>
        CreateMklOp;
  };

 private:
  std::unordered_map<std::string, MklEagerOp> mkl_eager_ops_;

  // The entry point to execute the op rewrite.
  Status Run(EagerOperation* orig_op,
             std::unique_ptr<tensorflow::EagerOperation>* out_op);

  // Initializes the new op and sets up its inputs and attributes
  static Status SetupNewOp(EagerOperation* orig_op, const string mkl_op_name,
                           std::unique_ptr<EagerOperation>* new_mkl_op);

  // Generic rewrite that can be used for any mkl op that doesn't need
  // special processing.
  static Status CreateGenericMklOp(EagerOperation* orig_op,
                                   std::unique_ptr<EagerOperation>* mkl_op);

  // Creates new MKL op for Conv2D, Conv2DBackpropInput and
  // Conv2DBackpropFilter.
  static Status CreateMklConv2DOp(
      EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_conv2d_op);

  // Rewrite rule for Conv2D, Conv2DBackpropInput and Conv2DBackpropFilter.
  static bool RewriteConv2D(EagerOperation* op);

  // Calls op-specific rewrite function to create new MKL op.
  Status RewriteToMklOp(EagerOperation* orig_op,
                        std::unique_ptr<EagerOperation>* mkl_op);

  // Check whether we can rewrite the op to MKL one or not.
  bool ShouldRewriteOp(EagerOperation* op);

  // Default rewrite rule to be used when rewrite should happen without any
  // restriction.
  static bool AlwaysRewrite(EagerOperation* op) { return true; }

  // Check if kernel is registered for a particular op.
  bool FastCheckIfKernelRegistered(std::string op_name, DataType dt);

  // This is called by FastCheckIfKernelRegistered once per unique op name
  // and data type.
  bool SlowCheckIfKernelRegistered(std::string op_name, DataType dt);

  // Helper function to insert mkl_eager_ops to Map
  void InsertMKLEagerOps(MklEagerOp op);

  // Map used by FastCheckIfKernelRegistered.
  std::unordered_map<std::string, bool> registered_kernels_map_;
};

REGISTER_REWRITE(EagerOpRewriteRegistry::PRE_EXECUTION, MklEagerOpRewrite);

// Constructor
MklEagerOpRewrite::MklEagerOpRewrite(string name, string file, string line)
    : EagerOpRewrite(name, file, line), registered_kernels_map_() {
  InsertMKLEagerOps({"BatchMatMul", AlwaysRewrite, CreateGenericMklOp});
  InsertMKLEagerOps({"BatchMatMulV2", AlwaysRewrite, CreateGenericMklOp});
  InsertMKLEagerOps({"Conv2D", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"Conv2DBackpropInput", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"Conv2DBackpropFilter", RewriteConv2D, CreateMklConv2DOp});
  InsertMKLEagerOps({"MatMul", AlwaysRewrite, CreateGenericMklOp});
};

void MklEagerOpRewrite::InsertMKLEagerOps(MklEagerOp op) {
  mkl_eager_ops_.insert(std::make_pair(op.op_name, op));
}

Status MklEagerOpRewrite::Run(
    EagerOperation* orig_op,
    std::unique_ptr<tensorflow::EagerOperation>* out_op) {
  if (ShouldRewriteOp(orig_op)) {
    TF_CHECK_OK(RewriteToMklOp(orig_op, out_op));
  }
  return Status::OK();
}

Status MklEagerOpRewrite::SetupNewOp(
    EagerOperation* orig_op, const string mkl_op_name,
    std::unique_ptr<EagerOperation>* new_mkl_op) {
  const tensorflow::AttrTypeMap* types;
  bool is_function = false;
  TF_RETURN_IF_ERROR(
      tensorflow::AttrTypeMapForOp(mkl_op_name.c_str(), &types, &is_function));
  EagerContext* ctx = orig_op->EagerContext();
  new_mkl_op->reset(new tensorflow::EagerOperation(ctx, mkl_op_name.c_str(),
                                                   is_function, types));

  int num_inputs = orig_op->Inputs().size();
  // Add all inputs to the new op.
  for (int i = 0; i < num_inputs; ++i) {
    (*new_mkl_op)->AddInput(orig_op->Inputs()[i]);
  }

  // Copy all attributes to the new op.
  const NodeDef& orig_ndef = orig_op->MutableAttrs()->BuildNodeDef();

  AttrSlice attr_list(orig_ndef);
  for (const auto& attr : attr_list) {
    (*new_mkl_op)->MutableAttrs()->Set(attr.first, attr.second);
  }

  (*new_mkl_op)
      ->MutableAttrs()
      ->Set("_kernel", mkl_op_registry::kMklNameChangeOpLabel);

  if (orig_op->Device() != nullptr) {
    (*new_mkl_op)->SetDevice(orig_op->Device());
  } else {
    string device_name = orig_op->GetDeviceName();
    (*new_mkl_op)->SetDeviceName(device_name.c_str());
  }
  return Status::OK();
}

Status MklEagerOpRewrite::CreateGenericMklOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_op) {
  const string mkl_op_name = mkl_op_registry::GetMklOpName(orig_op->Name());
  TF_CHECK_OK(SetupNewOp(orig_op, mkl_op_name, mkl_op));
  return Status::OK();
}

Status MklEagerOpRewrite::CreateMklConv2DOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_conv2d_op) {
  const string mkl_op_name =
      mkl_op_registry::GetMklEagerOpName(orig_op->Name());
  TF_CHECK_OK(SetupNewOp(orig_op, mkl_op_name, mkl_conv2d_op));
  return Status::OK();
}

bool MklEagerOpRewrite::ShouldRewriteOp(EagerOperation* op) {
  // Don't rewrite the op if MKL use is disabled at runtime.
  if (DisableMKL()) {
    return false;
  }
  DataType data_type;
  if (op->Attrs().Get("T", &data_type) != Status::OK()) {
    return false;
  }
  // Check if we have registered MKL kernel for this op.
  bool kernel_found = FastCheckIfKernelRegistered(op->Name(), data_type);
  if (!kernel_found) {
    return false;
  }

  // Find and call the op's rewrite rule that determines whether we need to
  // rewrite this op or not.
  auto it = mkl_eager_ops_.find(op->Name());
  if (it != mkl_eager_ops_.end()) {
    // Eager op found so verify Rewrite
    if (it->second.RewriteRule(op)) {
      return true;
    }
  }
  return false;
}

bool MklEagerOpRewrite::FastCheckIfKernelRegistered(std::string op_name,
                                                    DataType dt) {
  // Check for kernel registration only once per op name and data type
  // for performance reasons.
  string registered_kernels_key = op_name + std::to_string(dt);
  auto kernel_element = registered_kernels_map_.find(registered_kernels_key);
  bool kernel_registered = false;
  if (kernel_element == registered_kernels_map_.end()) {
    // Kernel registration is not verified even once yet.
    // So verify and store registration.
    kernel_registered = SlowCheckIfKernelRegistered(op_name, dt);
    registered_kernels_map_.insert(
        std::make_pair(registered_kernels_key, kernel_registered));
  } else {
    // Kernel is visited atleast once. return stored registration result.
    kernel_registered = kernel_element->second;
  }

  return kernel_registered;
}

bool MklEagerOpRewrite::SlowCheckIfKernelRegistered(string op_name,
                                                    DataType dt) {
  // Find if the eager op_name exists in mkl_eager_ops_ list.
  auto element = mkl_eager_ops_.find(op_name);
  if (element != mkl_eager_ops_.end() && dt == DT_FLOAT) {
    // Eager Op exists. So verify registry and return registered or not.
    return (mkl_op_registry::IsMklNameChangeOp(
                mkl_op_registry::GetMklEagerOpName(op_name), dt) ||
            mkl_op_registry::IsMklNameChangeOp(
                mkl_op_registry::GetMklOpName(op_name), dt));
  } else {
    return false;
  }
}

Status MklEagerOpRewrite::RewriteToMklOp(
    EagerOperation* orig_op, std::unique_ptr<EagerOperation>* mkl_op) {
  // TODO(intel-tf): mkl_eager_ops_ lookup can be reduced from twice
  // (once each in ShouldRewriteOp & RewriteToMklOp) to just once.
  mkl_eager_ops_[orig_op->Name()].CreateMklOp(orig_op, mkl_op);
  return Status::OK();
}

bool MklEagerOpRewrite::RewriteConv2D(EagerOperation* op) {
  const NodeDef& ndef = op->MutableAttrs()->BuildNodeDef();
  string padding;
  TF_CHECK_OK(GetNodeAttr(ndef, "padding", &padding));
  // Right now MKL Conv2D does not support explicit padding.
  return (padding != "EXPLICIT");
}

}  // namespace tensorflow
#endif  // INTEL_MKL
<|endoftext|>


<commit_before>/*M///////////////////////////////////////////////////////////////////////////////////////
//
//  IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
//
//  By downloading, copying, installing or using the software you agree to this license.
//  If you do not agree to this license, do not download, install,
//  copy or use the software.
//
//
//                           License Agreement
//                For Open Source Computer Vision Library
//
// Copyright (C) 2010-2012, Multicoreware, Inc., all rights reserved.
// Copyright (C) 2010-2012, Advanced Micro Devices, Inc., all rights reserved.
// Third party copyrights are property of their respective owners.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
//   * Redistribution's of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//
//   * Redistribution's in binary form must reproduce the above copyright notice,
//     this list of conditions and the following disclaimer in the documentation
//     and/or other materials provided with the distribution.
//
//   * The name of the copyright holders may not be used to endorse or promote products
//     derived from this software without specific prior written permission.
//
// This software is provided by the copyright holders and contributors "as is" and
// any express or implied warranties, including, but not limited to, the implied
// warranties of merchantability and fitness for a particular purpose are disclaimed.
// In no event shall the Intel Corporation or contributors be liable for any direct,
// indirect, incidental, special, exemplary, or consequential damages
// (including, but not limited to, procurement of substitute goods or services;
// loss of use, data, or profits; or business interruption) however caused
// and on any theory of liability, whether in contract, strict liability,
// or tort (including negligence or otherwise) arising in any way out of
// the use of this software, even if advised of the possibility of such damage.
//
//M*/

#include "perf_precomp.hpp"

#ifdef HAVE_OPENCL

#include "opencv2/ocl/ocl.hpp"
using namespace std;
using namespace testing;
using namespace perf;
using namespace cv;
using namespace cv::superres;

namespace
{
    class OneFrameSource_OCL : public FrameSource
    {
    public:
        explicit OneFrameSource_OCL(const ocl::oclMat& frame) : frame_(frame) {}

        void nextFrame(OutputArray frame)
        {
            ocl::getOclMatRef(frame) = frame_;
        }
        void reset()
        {
        }

    private:
        ocl::oclMat frame_;
    };

    class ZeroOpticalFlowOCL : public DenseOpticalFlowExt
    {
    public:
        void calc(InputArray frame0, InputArray, OutputArray flow1, OutputArray flow2)
        {
            ocl::oclMat& frame0_ = ocl::getOclMatRef(frame0);
            ocl::oclMat& flow1_ = ocl::getOclMatRef(flow1);
            ocl::oclMat& flow2_ = ocl::getOclMatRef(flow2);

            cv::Size size = frame0_.size();

            if(!flow2.needed())
            {
                flow1_.create(size, CV_32FC2);
                flow1_.setTo(Scalar::all(0));
            }
            else
            {
                flow1_.create(size, CV_32FC1);
                flow2_.create(size, CV_32FC1);

                flow1_.setTo(Scalar::all(0));
                flow2_.setTo(Scalar::all(0));
            }
        }

        void collectGarbage()
        {
        }
    };
}

PERF_TEST_P(Size_MatType, SuperResolution_BTVL1_OCL,
    Combine(Values(szSmall64, szSmall128),
    Values(MatType(CV_8UC1), MatType(CV_8UC3))))
{
    std::vector<cv::ocl::Info>info;
    cv::ocl::getDevice(info);

    declare.time(5 * 60);

    const Size size = get<0>(GetParam());
    const int type = get<1>(GetParam());

    Mat frame(size, type);
    declare.in(frame, WARMUP_RNG);

    ocl::oclMat frame_ocl;
    frame_ocl.upload(frame);

    const int scale = 2;
    const int iterations = 50;
    const int temporalAreaRadius = 1;
    Ptr<DenseOpticalFlowExt> opticalFlowOcl(new ZeroOpticalFlowOCL);

    Ptr<SuperResolution> superRes_ocl = createSuperResolution_BTVL1_OCL();

    superRes_ocl->set("scale", scale);
    superRes_ocl->set("iterations", iterations);
    superRes_ocl->set("temporalAreaRadius", temporalAreaRadius);
    superRes_ocl->set("opticalFlow", opticalFlowOcl);

    superRes_ocl->setInput(new OneFrameSource_OCL(frame_ocl));

    ocl::oclMat dst_ocl;
    superRes_ocl->nextFrame(dst_ocl);

    TEST_CYCLE_N(10) superRes_ocl->nextFrame(dst_ocl);
    frame_ocl.release();
    CPU_SANITY_CHECK(dst_ocl);
}
#endif
<commit_msg>Fixed a compilation error (get<N> not found).<commit_after>/*M///////////////////////////////////////////////////////////////////////////////////////
//
//  IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
//
//  By downloading, copying, installing or using the software you agree to this license.
//  If you do not agree to this license, do not download, install,
//  copy or use the software.
//
//
//                           License Agreement
//                For Open Source Computer Vision Library
//
// Copyright (C) 2010-2012, Multicoreware, Inc., all rights reserved.
// Copyright (C) 2010-2012, Advanced Micro Devices, Inc., all rights reserved.
// Third party copyrights are property of their respective owners.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
//   * Redistribution's of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//
//   * Redistribution's in binary form must reproduce the above copyright notice,
//     this list of conditions and the following disclaimer in the documentation
//     and/or other materials provided with the distribution.
//
//   * The name of the copyright holders may not be used to endorse or promote products
//     derived from this software without specific prior written permission.
//
// This software is provided by the copyright holders and contributors "as is" and
// any express or implied warranties, including, but not limited to, the implied
// warranties of merchantability and fitness for a particular purpose are disclaimed.
// In no event shall the Intel Corporation or contributors be liable for any direct,
// indirect, incidental, special, exemplary, or consequential damages
// (including, but not limited to, procurement of substitute goods or services;
// loss of use, data, or profits; or business interruption) however caused
// and on any theory of liability, whether in contract, strict liability,
// or tort (including negligence or otherwise) arising in any way out of
// the use of this software, even if advised of the possibility of such damage.
//
//M*/

#include "perf_precomp.hpp"

#ifdef HAVE_OPENCL

#include "opencv2/ocl/ocl.hpp"
using namespace std;
using namespace std::tr1;
using namespace testing;
using namespace perf;
using namespace cv;
using namespace cv::superres;

namespace
{
    class OneFrameSource_OCL : public FrameSource
    {
    public:
        explicit OneFrameSource_OCL(const ocl::oclMat& frame) : frame_(frame) {}

        void nextFrame(OutputArray frame)
        {
            ocl::getOclMatRef(frame) = frame_;
        }
        void reset()
        {
        }

    private:
        ocl::oclMat frame_;
    };

    class ZeroOpticalFlowOCL : public DenseOpticalFlowExt
    {
    public:
        void calc(InputArray frame0, InputArray, OutputArray flow1, OutputArray flow2)
        {
            ocl::oclMat& frame0_ = ocl::getOclMatRef(frame0);
            ocl::oclMat& flow1_ = ocl::getOclMatRef(flow1);
            ocl::oclMat& flow2_ = ocl::getOclMatRef(flow2);

            cv::Size size = frame0_.size();

            if(!flow2.needed())
            {
                flow1_.create(size, CV_32FC2);
                flow1_.setTo(Scalar::all(0));
            }
            else
            {
                flow1_.create(size, CV_32FC1);
                flow2_.create(size, CV_32FC1);

                flow1_.setTo(Scalar::all(0));
                flow2_.setTo(Scalar::all(0));
            }
        }

        void collectGarbage()
        {
        }
    };
}

PERF_TEST_P(Size_MatType, SuperResolution_BTVL1_OCL,
    Combine(Values(szSmall64, szSmall128),
    Values(MatType(CV_8UC1), MatType(CV_8UC3))))
{
    std::vector<cv::ocl::Info>info;
    cv::ocl::getDevice(info);

    declare.time(5 * 60);

    const Size size = get<0>(GetParam());
    const int type = get<1>(GetParam());

    Mat frame(size, type);
    declare.in(frame, WARMUP_RNG);

    ocl::oclMat frame_ocl;
    frame_ocl.upload(frame);

    const int scale = 2;
    const int iterations = 50;
    const int temporalAreaRadius = 1;
    Ptr<DenseOpticalFlowExt> opticalFlowOcl(new ZeroOpticalFlowOCL);

    Ptr<SuperResolution> superRes_ocl = createSuperResolution_BTVL1_OCL();

    superRes_ocl->set("scale", scale);
    superRes_ocl->set("iterations", iterations);
    superRes_ocl->set("temporalAreaRadius", temporalAreaRadius);
    superRes_ocl->set("opticalFlow", opticalFlowOcl);

    superRes_ocl->setInput(new OneFrameSource_OCL(frame_ocl));

    ocl::oclMat dst_ocl;
    superRes_ocl->nextFrame(dst_ocl);

    TEST_CYCLE_N(10) superRes_ocl->nextFrame(dst_ocl);
    frame_ocl.release();
    CPU_SANITY_CHECK(dst_ocl);
}
#endif
<|endoftext|>


<commit_before>// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <iostream>

#include "base/at_exit.h"
#include "base/command_line.h"
#include "base/file_path.h"
#include "base/logging.h"
#include "base/path_service.h"
#include "base/process_util.h"
#include "base/string_number_conversions.h"
#include "base/string_piece.h"
#include "base/string_util.h"
#include "base/utf_string_conversions.h"
#include "base/memory/scoped_ptr.h"
#include "base/threading/thread.h"
#include "sawbuck/common/com_utils.h"
#include "syzygy/trace/protocol/call_trace_defs.h"
#include "syzygy/trace/rpc/rpc_helpers.h"
#include "syzygy/trace/service/service.h"
#include "syzygy/trace/service/service_rpc_impl.h"
#include "syzygy/trace/service/trace_file_writer_factory.h"

namespace trace {
namespace service {
namespace {

using ::trace::client::CreateRpcBinding;
using ::trace::client::InvokeRpc;

// Minimum buffer size to allow (1 MB).
const int kMinBufferSize = 1024 * 1024;

// Minumum number of buffers to allocate.
const int kMinBuffers = 16;

// A static location to which the current instance id can be saved. We
// persist it here so that OnConsoleCtrl can have access to the instance
// id when it is invoked on the signal handler thread.
wchar_t saved_instance_id[16] = {0};

// Forward declaration.
bool StopService(const base::StringPiece16& instance_id);

// Handler function to be called on exit signals (Ctrl-C, TERM, etc...).
BOOL WINAPI OnConsoleCtrl(DWORD ctrl_type) {
  if (ctrl_type != CTRL_LOGOFF_EVENT) {
    StopService(saved_instance_id);
    return TRUE;
  }
  return FALSE;
}

const char* const kInstanceId = "instance-id";

const char kUsage[] =
    "Usage: call_trace_service ACTION [OPTIONS]\n"
    "\n"
    "Actions:\n"
    "  start              Start the call trace service. This causes an\n"
    "                     instance of the service to be launched as a\n"
    "                     foreground process.\n"
    "  spawn              Spawns an instance of the call trace service, waits\n"
    "                     for it to be ready, and returns. The call trace\n"
    "                     service continues running in the background.\n"
    "  stop               Stop the call trace service.\n"
    "\n"
    "Options:\n"
    "  --help             Show this help message.\n"
    "  --trace-dir=PATH   The directory in which to write the trace files.\n"
    "  --buffer-size=NUM  The size (in bytes) of each buffer to allocate.\n"
    "  --num-incremental-buffers=NUM\n"
    "                     The number of buffers by which to grow the buffer\n"
    "                     pool each time the client exhausts its available\n"
    "                     buffer space.\n"
    "  --enable-exits     Enable exit tracing (off by default).\n"
    "  --verbose          Increase the logging verbosity to also include\n"
    "                     debug-level information.\n"
    "  --instance-id=ID   A unique identifier to use for the RPC endoint.\n"
    "                     This allows multiple instances of the service to\n"
    "                     run concurrently. By default this is empty.\n"
    "\n";

int Usage() {
  std::cout << kUsage;
  return 1;
}

bool GetInstanceId(const CommandLine* cmd_line, std::wstring* id) {
  DCHECK(cmd_line != NULL);
  DCHECK(id != NULL);

  // If not specified, this defaults to the empty string.
  *id = cmd_line->GetSwitchValueNative(kInstanceId);

  const size_t kMaxLength = arraysize(saved_instance_id) - 1;
  if (id->length() > kMaxLength) {
    LOG(ERROR) << "The instance id '" << *id << "' is too long. "
               << "The max length is " << kMaxLength << " characters.";
    return false;
  }

  return true;
}

bool RunService(const CommandLine* cmd_line) {
  DCHECK(cmd_line != NULL);

  base::Thread writer_thread("trace-file-writer");
  if (!writer_thread.StartWithOptions(
          base::Thread::Options(MessageLoop::TYPE_IO, 0))) {
    LOG(ERROR) << "Failed to start call trace service writer thread.";
    return 1;
  }

  MessageLoop* message_loop = writer_thread.message_loop();
  TraceFileWriterFactory trace_file_writer_factory(message_loop);
  Service call_trace_service(&trace_file_writer_factory);
  RpcServiceInstanceManager rpc_instance(&call_trace_service);

  // Get/set the instance id.
  std::wstring instance_id;
  if (!GetInstanceId(cmd_line, &instance_id))
    return false;

  call_trace_service.set_instance_id(instance_id);
  base::wcslcpy(saved_instance_id,
                instance_id.c_str(),
                arraysize(saved_instance_id));

  // Set up the trace directory.
  FilePath trace_directory(cmd_line->GetSwitchValuePath("trace-dir"));
  if (trace_directory.empty())
    trace_directory = FilePath(L".");
  if (!trace_file_writer_factory.SetTraceFileDirectory(trace_directory))
    return false;

  // Setup the buffer size.
  std::wstring buffer_size_str(cmd_line->GetSwitchValueNative("buffer-size"));
  if (!buffer_size_str.empty()) {
    int num = 0;
    if (!base::StringToInt(buffer_size_str, &num) || num < kMinBufferSize) {
      LOG(ERROR) << "Buffer size is too small (<" << kMinBufferSize << ").";
      return false;
    }
    call_trace_service.set_buffer_size_in_bytes(num);
  }

  if (cmd_line->HasSwitch("enable-exits")) {
    call_trace_service.set_flags(TRACE_FLAG_ENTER | TRACE_FLAG_EXIT);
  }

  // Setup the number of incremental buffers
  std::wstring buffers_str(
      cmd_line->GetSwitchValueNative("num-incremental-buffers"));
  if (!buffers_str.empty()) {
    int num = 0;
    if (!base::StringToInt(buffers_str, &num) || num < kMinBuffers) {
      LOG(ERROR) << "Number of incremental buffers is too small (<"
                 << kMinBuffers << ").";
      return false;
    }
    call_trace_service.set_num_incremental_buffers(num);
  }

  // Setup the handler for exit signals.
  if (!SetConsoleCtrlHandler(&OnConsoleCtrl, TRUE)) {
    DWORD error = ::GetLastError();
    LOG(ERROR) << "Failed to register shutdown handler: "
               << com::LogWe(error) << ".";
    return false;
  }

  // Run the service in blocking mode. This will not return until the service
  // has been externally stopped.
  call_trace_service.Start(false);

  // We no longer need to look out for exit signals.
  SetConsoleCtrlHandler(&OnConsoleCtrl, FALSE);

  // The call trace service will be stopped on destruction.
  return true;
}

bool SpawnService(const CommandLine* cmd_line) {
  // Get the path to ourselves.
  FilePath self_path;
  PathService::Get(base::FILE_EXE, &self_path);

  // Build a command line for starting a new instance of the service.
  CommandLine service_cmd(self_path);
  service_cmd.AppendArg("start");

  // Copy over any other switches.
  CommandLine::SwitchMap::const_iterator it =
      service_cmd.GetSwitches().begin();
  for (; it != service_cmd.GetSwitches().end(); ++it)
    service_cmd.AppendSwitchNative(it->first, it->second);

  // Get the instance id.
  std::wstring instance_id;
  if (!GetInstanceId(cmd_line, &instance_id))
    return false;

  // Launch a new process in the background.
  LOG(INFO) << "Launching background call trace service with instance ID \""
            << instance_id << "\".";
  base::ProcessHandle service_process;
  base::LaunchOptions options;
  options.start_hidden = true;
  if (!base::LaunchProcess(service_cmd, options, &service_process)) {
    LOG(ERROR) << "Failed to launch process.";
    return false;
  }
  DCHECK_NE(base::kNullProcessHandle, service_process);

  // Get the name of the event that will be signalled when the service is up
  // and running.
  std::wstring event_name;
  ::GetSyzygyCallTraceRpcEventName(instance_id, &event_name);
  base::win::ScopedHandle rpc_event(
      ::CreateEvent(NULL, TRUE, FALSE, event_name.c_str()));
  if (!rpc_event.IsValid()) {
    LOG(ERROR) << "Unable to create RPC event for instance id \""
               << instance_id << "\".";
    return false;
  }

  // We wait on both the RPC event and the process, as if the process fails for
  // any reason, it'll exit and its handle will become signalled.
  HANDLE handles[] = { rpc_event.Get(), service_process };
  if (::WaitForMultipleObjects(arraysize(handles),
                               handles,
                               FALSE,
                               INFINITE) != WAIT_OBJECT_0) {
    LOG(ERROR) << "The spawned call trace service exited in error.";
    return false;
  }

  LOG(INFO) << "Background call trace service with instance ID \""
            << instance_id << "\" is ready.";

  return true;
}

bool StopService(const base::StringPiece16& instance_id) {
  std::wstring protocol;
  std::wstring endpoint;

  ::GetSyzygyCallTraceRpcProtocol(&protocol);
  ::GetSyzygyCallTraceRpcEndpoint(instance_id, &endpoint);

  LOG(INFO) << "Stopping call trace logging service instance at '"
            << endpoint << "' via " << protocol << '.';

  handle_t binding = NULL;
  if (!CreateRpcBinding(protocol, endpoint, &binding)) {
    LOG(ERROR) << "Failed to connect to call trace logging service.";
    return false;
  }

  if (!InvokeRpc(CallTraceClient_Stop, binding).succeeded()) {
    LOG(ERROR) << "Failed to stop call trace logging service.";
    return false;
  }

  // TODO(rogerm): It would be nice to make this blocking until the
  //    service actually shuts down. Perhaps with retries on the stop
  //    request.
  LOG(INFO) << "Call trace service shutdown has been requested.";
  return true;
}

extern "C" int main(int argc, char** argv) {
  base::AtExitManager at_exit_manager;
  CommandLine::Init(argc, argv);
  const int kVlogLevelVerbose = -2;
  if (!logging::InitLogging(
          L"",
          logging::LOG_ONLY_TO_SYSTEM_DEBUG_LOG,
          logging::DONT_LOCK_LOG_FILE,
          logging::APPEND_TO_OLD_LOG_FILE,
          logging::ENABLE_DCHECK_FOR_NON_OFFICIAL_RELEASE_BUILDS)) {
    return 1;
  }

  // TODO(rogerm): Turn on ETW logging as well.

  CommandLine* cmd_line = CommandLine::ForCurrentProcess();
  DCHECK(cmd_line != NULL);

  if (cmd_line->HasSwitch("verbose")) {
    logging::SetMinLogLevel(kVlogLevelVerbose);
  }

  if (cmd_line->HasSwitch("help") || cmd_line->GetArgs().size() < 1) {
    return Usage();
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "stop")) {
    std::wstring id;
    return (GetInstanceId(cmd_line, &id) && StopService(id)) ? 0 : 1;
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "start")) {
    return RunService(cmd_line) ? 0 : 1;
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "spawn")) {
    return SpawnService(cmd_line) ? 0 : 1;
  }

  return Usage();
}

}  // namespace
}  // namespace service
}  // namespace trace
<commit_msg>Fix a boneheaded bug in 'spawn' functionality.<commit_after>// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <iostream>

#include "base/at_exit.h"
#include "base/command_line.h"
#include "base/file_path.h"
#include "base/logging.h"
#include "base/path_service.h"
#include "base/process_util.h"
#include "base/string_number_conversions.h"
#include "base/string_piece.h"
#include "base/string_util.h"
#include "base/utf_string_conversions.h"
#include "base/memory/scoped_ptr.h"
#include "base/threading/thread.h"
#include "sawbuck/common/com_utils.h"
#include "syzygy/trace/protocol/call_trace_defs.h"
#include "syzygy/trace/rpc/rpc_helpers.h"
#include "syzygy/trace/service/service.h"
#include "syzygy/trace/service/service_rpc_impl.h"
#include "syzygy/trace/service/trace_file_writer_factory.h"

namespace trace {
namespace service {
namespace {

using ::trace::client::CreateRpcBinding;
using ::trace::client::InvokeRpc;

// Minimum buffer size to allow (1 MB).
const int kMinBufferSize = 1024 * 1024;

// Minumum number of buffers to allocate.
const int kMinBuffers = 16;

// A static location to which the current instance id can be saved. We
// persist it here so that OnConsoleCtrl can have access to the instance
// id when it is invoked on the signal handler thread.
wchar_t saved_instance_id[16] = {0};

// Forward declaration.
bool StopService(const base::StringPiece16& instance_id);

// Handler function to be called on exit signals (Ctrl-C, TERM, etc...).
BOOL WINAPI OnConsoleCtrl(DWORD ctrl_type) {
  if (ctrl_type != CTRL_LOGOFF_EVENT) {
    StopService(saved_instance_id);
    return TRUE;
  }
  return FALSE;
}

const char* const kInstanceId = "instance-id";

const char kUsage[] =
    "Usage: call_trace_service ACTION [OPTIONS]\n"
    "\n"
    "Actions:\n"
    "  start              Start the call trace service. This causes an\n"
    "                     instance of the service to be launched as a\n"
    "                     foreground process.\n"
    "  spawn              Spawns an instance of the call trace service, waits\n"
    "                     for it to be ready, and returns. The call trace\n"
    "                     service continues running in the background.\n"
    "  stop               Stop the call trace service.\n"
    "\n"
    "Options:\n"
    "  --help             Show this help message.\n"
    "  --trace-dir=PATH   The directory in which to write the trace files.\n"
    "  --buffer-size=NUM  The size (in bytes) of each buffer to allocate.\n"
    "  --num-incremental-buffers=NUM\n"
    "                     The number of buffers by which to grow the buffer\n"
    "                     pool each time the client exhausts its available\n"
    "                     buffer space.\n"
    "  --enable-exits     Enable exit tracing (off by default).\n"
    "  --verbose          Increase the logging verbosity to also include\n"
    "                     debug-level information.\n"
    "  --instance-id=ID   A unique identifier to use for the RPC endoint.\n"
    "                     This allows multiple instances of the service to\n"
    "                     run concurrently. By default this is empty.\n"
    "\n";

int Usage() {
  std::cout << kUsage;
  return 1;
}

bool GetInstanceId(const CommandLine* cmd_line, std::wstring* id) {
  DCHECK(cmd_line != NULL);
  DCHECK(id != NULL);

  // If not specified, this defaults to the empty string.
  *id = cmd_line->GetSwitchValueNative(kInstanceId);

  const size_t kMaxLength = arraysize(saved_instance_id) - 1;
  if (id->length() > kMaxLength) {
    LOG(ERROR) << "The instance id '" << *id << "' is too long. "
               << "The max length is " << kMaxLength << " characters.";
    return false;
  }

  return true;
}

bool RunService(const CommandLine* cmd_line) {
  DCHECK(cmd_line != NULL);

  base::Thread writer_thread("trace-file-writer");
  if (!writer_thread.StartWithOptions(
          base::Thread::Options(MessageLoop::TYPE_IO, 0))) {
    LOG(ERROR) << "Failed to start call trace service writer thread.";
    return 1;
  }

  MessageLoop* message_loop = writer_thread.message_loop();
  TraceFileWriterFactory trace_file_writer_factory(message_loop);
  Service call_trace_service(&trace_file_writer_factory);
  RpcServiceInstanceManager rpc_instance(&call_trace_service);

  // Get/set the instance id.
  std::wstring instance_id;
  if (!GetInstanceId(cmd_line, &instance_id))
    return false;

  call_trace_service.set_instance_id(instance_id);
  base::wcslcpy(saved_instance_id,
                instance_id.c_str(),
                arraysize(saved_instance_id));

  // Set up the trace directory.
  FilePath trace_directory(cmd_line->GetSwitchValuePath("trace-dir"));
  if (trace_directory.empty())
    trace_directory = FilePath(L".");
  if (!trace_file_writer_factory.SetTraceFileDirectory(trace_directory))
    return false;

  // Setup the buffer size.
  std::wstring buffer_size_str(cmd_line->GetSwitchValueNative("buffer-size"));
  if (!buffer_size_str.empty()) {
    int num = 0;
    if (!base::StringToInt(buffer_size_str, &num) || num < kMinBufferSize) {
      LOG(ERROR) << "Buffer size is too small (<" << kMinBufferSize << ").";
      return false;
    }
    call_trace_service.set_buffer_size_in_bytes(num);
  }

  if (cmd_line->HasSwitch("enable-exits")) {
    call_trace_service.set_flags(TRACE_FLAG_ENTER | TRACE_FLAG_EXIT);
  }

  // Setup the number of incremental buffers
  std::wstring buffers_str(
      cmd_line->GetSwitchValueNative("num-incremental-buffers"));
  if (!buffers_str.empty()) {
    int num = 0;
    if (!base::StringToInt(buffers_str, &num) || num < kMinBuffers) {
      LOG(ERROR) << "Number of incremental buffers is too small (<"
                 << kMinBuffers << ").";
      return false;
    }
    call_trace_service.set_num_incremental_buffers(num);
  }

  // Setup the handler for exit signals.
  if (!SetConsoleCtrlHandler(&OnConsoleCtrl, TRUE)) {
    DWORD error = ::GetLastError();
    LOG(ERROR) << "Failed to register shutdown handler: "
               << com::LogWe(error) << ".";
    return false;
  }

  // Run the service in blocking mode. This will not return until the service
  // has been externally stopped.
  call_trace_service.Start(false);

  // We no longer need to look out for exit signals.
  SetConsoleCtrlHandler(&OnConsoleCtrl, FALSE);

  // The call trace service will be stopped on destruction.
  return true;
}

bool SpawnService(const CommandLine* cmd_line) {
  // Get the path to ourselves.
  FilePath self_path;
  PathService::Get(base::FILE_EXE, &self_path);

  // Build a command line for starting a new instance of the service.
  CommandLine service_cmd(self_path);
  service_cmd.AppendArg("start");

  // Copy over any other switches.
  CommandLine::SwitchMap::const_iterator it =
      cmd_line->GetSwitches().begin();
  for (; it != cmd_line->GetSwitches().end(); ++it)
    service_cmd.AppendSwitchNative(it->first, it->second);

  // Get the instance id.
  std::wstring instance_id;
  if (!GetInstanceId(cmd_line, &instance_id))
    return false;

  // Launch a new process in the background.
  LOG(INFO) << "Launching background call trace service with instance ID \""
            << instance_id << "\".";
  base::ProcessHandle service_process;
  base::LaunchOptions options;
  options.start_hidden = true;
  if (!base::LaunchProcess(service_cmd, options, &service_process)) {
    LOG(ERROR) << "Failed to launch process.";
    return false;
  }
  DCHECK_NE(base::kNullProcessHandle, service_process);

  // Get the name of the event that will be signalled when the service is up
  // and running.
  std::wstring event_name;
  ::GetSyzygyCallTraceRpcEventName(instance_id, &event_name);
  base::win::ScopedHandle rpc_event(
      ::CreateEvent(NULL, TRUE, FALSE, event_name.c_str()));
  if (!rpc_event.IsValid()) {
    LOG(ERROR) << "Unable to create RPC event for instance id \""
               << instance_id << "\".";
    return false;
  }

  // We wait on both the RPC event and the process, as if the process fails for
  // any reason, it'll exit and its handle will become signalled.
  HANDLE handles[] = { rpc_event.Get(), service_process };
  if (::WaitForMultipleObjects(arraysize(handles),
                               handles,
                               FALSE,
                               INFINITE) != WAIT_OBJECT_0) {
    LOG(ERROR) << "The spawned call trace service exited in error.";
    return false;
  }

  LOG(INFO) << "Background call trace service with instance ID \""
            << instance_id << "\" is ready.";

  return true;
}

bool StopService(const base::StringPiece16& instance_id) {
  std::wstring protocol;
  std::wstring endpoint;

  ::GetSyzygyCallTraceRpcProtocol(&protocol);
  ::GetSyzygyCallTraceRpcEndpoint(instance_id, &endpoint);

  LOG(INFO) << "Stopping call trace logging service instance at '"
            << endpoint << "' via " << protocol << '.';

  handle_t binding = NULL;
  if (!CreateRpcBinding(protocol, endpoint, &binding)) {
    LOG(ERROR) << "Failed to connect to call trace logging service.";
    return false;
  }

  if (!InvokeRpc(CallTraceClient_Stop, binding).succeeded()) {
    LOG(ERROR) << "Failed to stop call trace logging service.";
    return false;
  }

  // TODO(rogerm): It would be nice to make this blocking until the
  //    service actually shuts down. Perhaps with retries on the stop
  //    request.
  LOG(INFO) << "Call trace service shutdown has been requested.";
  return true;
}

extern "C" int main(int argc, char** argv) {
  base::AtExitManager at_exit_manager;
  CommandLine::Init(argc, argv);
  const int kVlogLevelVerbose = -2;
  if (!logging::InitLogging(
          L"",
          logging::LOG_ONLY_TO_SYSTEM_DEBUG_LOG,
          logging::DONT_LOCK_LOG_FILE,
          logging::APPEND_TO_OLD_LOG_FILE,
          logging::ENABLE_DCHECK_FOR_NON_OFFICIAL_RELEASE_BUILDS)) {
    return 1;
  }

  // TODO(rogerm): Turn on ETW logging as well.

  CommandLine* cmd_line = CommandLine::ForCurrentProcess();
  DCHECK(cmd_line != NULL);

  if (cmd_line->HasSwitch("verbose")) {
    logging::SetMinLogLevel(kVlogLevelVerbose);
  }

  if (cmd_line->HasSwitch("help") || cmd_line->GetArgs().size() < 1) {
    return Usage();
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "stop")) {
    std::wstring id;
    return (GetInstanceId(cmd_line, &id) && StopService(id)) ? 0 : 1;
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "start")) {
    return RunService(cmd_line) ? 0 : 1;
  }

  if (LowerCaseEqualsASCII(cmd_line->GetArgs()[0], "spawn")) {
    return SpawnService(cmd_line) ? 0 : 1;
  }

  return Usage();
}

}  // namespace
}  // namespace service
}  // namespace trace
<|endoftext|>