<commit_before>#include <mfapi.h>
//#include <mfobjects.h>
#include <mfidl.h>
#include <mferror.h>
#include "webmmfbytestreamhandler.hpp"
#include "webmmfsource.hpp"
#include <new>
#include <cassert>
#ifdef _DEBUG
#include "odbgstream.hpp"
using std::endl;
#endif

_COM_SMARTPTR_TYPEDEF(IMFMediaSource, __uuidof(IMFMediaSource));

namespace WebmMfSourceLib
{

HRESULT CreateHandler(
    IClassFactory* pClassFactory,
    IUnknown* pOuter,
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    *ppv = 0;

    if (pOuter)
        return CLASS_E_NOAGGREGATION;

    typedef WebmMfByteStreamHandler H;

    H* const p = new (std::nothrow) H(pClassFactory);

    if (p == 0)
        return E_OUTOFMEMORY;

    IUnknown* const pUnk = p;

    const HRESULT hr = pUnk->QueryInterface(iid, ppv);

    const ULONG cRef = pUnk->Release();
    cRef;

    return hr;
}

#pragma warning(disable:4355)  //'this' ptr in member init list
WebmMfByteStreamHandler::WebmMfByteStreamHandler(
    IClassFactory* pClassFactory) :
    m_pClassFactory(pClassFactory),
    m_cRef(1),
    m_bCancel(FALSE),
    m_async_load(this)
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler: ctor: this=0x"
       << (const void*)this
       << endl;
#endif

    HRESULT hr = m_pClassFactory->LockServer(TRUE);
    assert(SUCCEEDED(hr));

    hr = CLockable::Init();
    assert(SUCCEEDED(hr));
}
#pragma warning(default:4355)

WebmMfByteStreamHandler::~WebmMfByteStreamHandler()
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler: dtor; this=0x"
       << (const void*)this
       << endl;
#endif

    const HRESULT hr = m_pClassFactory->LockServer(FALSE);
    assert(SUCCEEDED(hr));
}

HRESULT WebmMfByteStreamHandler::QueryInterface(
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);

    if (iid == __uuidof(IUnknown))
    {
        pUnk = this;  //must be nondelegating
    }
    else if (iid == __uuidof(IMFByteStreamHandler))
    {
        pUnk = static_cast<IMFByteStreamHandler*>(this);
    }
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }

    pUnk->AddRef();
    return S_OK;
}

ULONG WebmMfByteStreamHandler::AddRef()
{
    const ULONG n = InterlockedIncrement(&m_cRef);

#if 0 //def _DEBUG
    odbgstream os;
    os << "WebmMfByteStreamHandler::AddRef: n=" << n << endl;
#endif

    return n;
}

ULONG WebmMfByteStreamHandler::Release()
{
    const LONG n = InterlockedDecrement(&m_cRef);

#if 0 //def _DEBUG
    odbgstream os;
    os << "WebmMfByteStreamHandler::Release: n=" << n << endl;
#endif

    if (n == 0)
        delete this;

    return n;
}

HRESULT WebmMfByteStreamHandler::BeginCreateObject(
    IMFByteStream* pByteStream,
    LPCWSTR pURL,
    DWORD dwFlags,
    IPropertyStore*,
    IUnknown** ppCancelCookie,
    IMFAsyncCallback* pCallback,
    IUnknown* pState)
{
    pURL;  //?

#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::BeginCreateObject: url="
       << (pURL ? pURL : L"<NONE>")
       << endl;
#endif

    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

    if (ppCancelCookie)
        *ppCancelCookie = 0;

    if (pByteStream == 0)
        return E_INVALIDARG;

    if (pCallback == 0)
        return E_INVALIDARG;

    if ((dwFlags & MF_RESOLUTION_MEDIASOURCE) == 0)
        return E_INVALIDARG;

    if (m_pResult)  //assume one-at-a-time creation
        return MF_E_UNEXPECTED;

    DWORD dw;
    hr = pByteStream->GetCapabilities(&dw);

    if (SUCCEEDED(hr))
    {
#ifdef _DEBUG
        odbgstream os;
        os << "webmmfsource::BeginCreateObject: caps:";

        if (dw & MFBYTESTREAM_IS_READABLE)
            os << " READABLE";

        if (dw & MFBYTESTREAM_IS_WRITABLE)
            os << " WRITABLE";

        if (dw & MFBYTESTREAM_IS_SEEKABLE)
            os << " SEEKABLE";

        if (dw & MFBYTESTREAM_IS_REMOTE)
            os << " REMOTE";

        if (dw & MFBYTESTREAM_IS_DIRECTORY)
            os << " DIRECTORY";

        if (dw & MFBYTESTREAM_HAS_SLOW_SEEK)
            os << " SLOW_SEEK";

        if (dw & MFBYTESTREAM_IS_PARTIALLY_DOWNLOADED)
            os << " PARTIALLY_DOWNLOADED";

        if (dw & MFBYTESTREAM_SHARE_WRITE)
            os << " SHARE_WRITE";

        os << endl;
#endif

        __noop;  //TODO: vet capabilities
    }

    WebmMfSource* pSource;

    hr = WebmMfSource::CreateSource(m_pClassFactory, pByteStream, pSource);

    if (FAILED(hr))
        return hr;

    const IMFMediaSourcePtr pUnk(pSource, false);

    IMFAsyncResultPtr pResult;

    hr = MFCreateAsyncResult(pUnk, pCallback, pState, &pResult);

    if (FAILED(hr))
        return hr;

    hr = pSource->BeginLoad(&m_async_load);

    if (FAILED(hr))
        return hr;

    m_pResult = pResult;

    m_bCancel = FALSE;
    m_pByteStream = pByteStream;

    return S_OK;
}

HRESULT WebmMfByteStreamHandler::EndCreateObject(
    IMFAsyncResult* pResult,
    MF_OBJECT_TYPE* pObjectType,
    IUnknown** ppObject)
{
    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::EndCreateObject (begin)" << endl;
#endif

    if (ppObject == 0)
        return E_POINTER;

    IUnknown*& pObject = *ppObject;
    pObject = 0;

    if (pObjectType == 0)
        return E_POINTER;

    MF_OBJECT_TYPE& type = *pObjectType;
    type = MF_OBJECT_INVALID;

    if (pResult == 0)
        return E_INVALIDARG;

    hr = pResult->GetStatus();

    if (FAILED(hr))  //for example, cancelled
        return hr;

    IUnknownPtr pUnk;

    hr = pResult->GetObject(&pUnk);

    if (FAILED(hr))  //shouldn't happen
        return hr;

    assert(pUnk);

    hr = pUnk->QueryInterface(__uuidof(IMFMediaSource), (void**)&pObject);

    if (FAILED(hr))
        return hr;

    type = MF_OBJECT_MEDIASOURCE;

#ifdef _DEBUG
    os << L"WebmMfByteStreamHandler::EndCreateObject (end)" << endl;
#endif

    return S_OK;
}

HRESULT WebmMfByteStreamHandler::CancelObjectCreation(IUnknown*)
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::CancelObjectCreation" << endl;
#endif

    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

    if (m_bCancel)
        return S_OK;

    if (!m_pByteStream)
        return S_OK;

    m_bCancel = TRUE;

    return m_pByteStream->Close();
}

HRESULT WebmMfByteStreamHandler::GetMaxNumberOfBytesRequiredForResolution(
    QWORD* pcb)
{
    if (pcb == 0)
        return E_POINTER;

    *pcb = 1024;  //TODO: ask the webm parser for this value

    return S_OK;
}

WebmMfByteStreamHandler::CAsyncLoad::CAsyncLoad(
    WebmMfByteStreamHandler* p) :
    m_pHandler(p)
{
}

WebmMfByteStreamHandler::CAsyncLoad::~CAsyncLoad()
{
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::QueryInterface(
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);

    if (iid == __uuidof(IUnknown))
    {
        pUnk = static_cast<IMFAsyncCallback*>(this);  //must be nondelegating
    }
    else if (iid == __uuidof(IMFAsyncCallback))
    {
        pUnk = static_cast<IMFAsyncCallback*>(this);
    }
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }

    pUnk->AddRef();
    return S_OK;
}

ULONG WebmMfByteStreamHandler::CAsyncLoad::AddRef()
{
    return m_pHandler->AddRef();
}

ULONG WebmMfByteStreamHandler::CAsyncLoad::Release()
{
    return m_pHandler->Release();
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::GetParameters(
    DWORD*,
    DWORD*)
{
    return E_NOTIMPL;  //means "assume default behavior"
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::Invoke(IMFAsyncResult* pResult)
{
    if (pResult == 0)
        return E_INVALIDARG;

    Lock lock;

    HRESULT hr = lock.Seize(m_pHandler);

    if (FAILED(hr))
        return hr;

    return m_pHandler->EndLoad(pResult);
}

HRESULT WebmMfByteStreamHandler::EndLoad(IMFAsyncResult* pLoadResult)
{
    if (!m_pResult)  //should never happen
        return MF_E_UNEXPECTED;

    HRESULT hrLoad;

    if (m_bCancel)
        hrLoad = MF_E_OPERATION_CANCELLED;
    else
        hrLoad = pLoadResult->GetStatus();

    HRESULT hr = m_pResult->SetStatus(hrLoad);
    assert(SUCCEEDED(hr));

    if (FAILED(hrLoad))
    {
        IUnknownPtr pUnk;

        hr = m_pResult->GetObject(&pUnk);

        if (SUCCEEDED(hr))
        {
            const IMFMediaSourcePtr pSource(pUnk);
            assert(pSource);

            hr = pSource->Shutdown();
        }
    }

    hr = MFInvokeCallback(m_pResult);

    m_pByteStream = 0;
    m_pResult = 0;

    return hr;
}

}  //end namespace WebmMfSourceLib
<commit_msg>webmmfsource: check byte stream length during creation<commit_after>#include <mfapi.h>
//#include <mfobjects.h>
#include <mfidl.h>
#include <mferror.h>
#include "webmmfbytestreamhandler.hpp"
#include "webmmfsource.hpp"
#include <new>
#include <cassert>
#ifdef _DEBUG
#include "odbgstream.hpp"
using std::endl;
#endif

_COM_SMARTPTR_TYPEDEF(IMFMediaSource, __uuidof(IMFMediaSource));

namespace WebmMfSourceLib
{

HRESULT CreateHandler(
    IClassFactory* pClassFactory,
    IUnknown* pOuter,
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    *ppv = 0;

    if (pOuter)
        return CLASS_E_NOAGGREGATION;

    typedef WebmMfByteStreamHandler H;

    H* const p = new (std::nothrow) H(pClassFactory);

    if (p == 0)
        return E_OUTOFMEMORY;

    IUnknown* const pUnk = p;

    const HRESULT hr = pUnk->QueryInterface(iid, ppv);

    const ULONG cRef = pUnk->Release();
    cRef;

    return hr;
}

#pragma warning(disable:4355)  //'this' ptr in member init list
WebmMfByteStreamHandler::WebmMfByteStreamHandler(
    IClassFactory* pClassFactory) :
    m_pClassFactory(pClassFactory),
    m_cRef(1),
    m_bCancel(FALSE),
    m_async_load(this)
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler: ctor: this=0x"
       << (const void*)this
       << endl;
#endif

    HRESULT hr = m_pClassFactory->LockServer(TRUE);
    assert(SUCCEEDED(hr));

    hr = CLockable::Init();
    assert(SUCCEEDED(hr));
}
#pragma warning(default:4355)

WebmMfByteStreamHandler::~WebmMfByteStreamHandler()
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler: dtor; this=0x"
       << (const void*)this
       << endl;
#endif

    const HRESULT hr = m_pClassFactory->LockServer(FALSE);
    assert(SUCCEEDED(hr));
}

HRESULT WebmMfByteStreamHandler::QueryInterface(
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);

    if (iid == __uuidof(IUnknown))
    {
        pUnk = this;  //must be nondelegating
    }
    else if (iid == __uuidof(IMFByteStreamHandler))
    {
        pUnk = static_cast<IMFByteStreamHandler*>(this);
    }
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }

    pUnk->AddRef();
    return S_OK;
}

ULONG WebmMfByteStreamHandler::AddRef()
{
    const ULONG n = InterlockedIncrement(&m_cRef);

#if 0 //def _DEBUG
    odbgstream os;
    os << "WebmMfByteStreamHandler::AddRef: n=" << n << endl;
#endif

    return n;
}

ULONG WebmMfByteStreamHandler::Release()
{
    const LONG n = InterlockedDecrement(&m_cRef);

#if 0 //def _DEBUG
    odbgstream os;
    os << "WebmMfByteStreamHandler::Release: n=" << n << endl;
#endif

    if (n == 0)
        delete this;

    return n;
}

HRESULT WebmMfByteStreamHandler::BeginCreateObject(
    IMFByteStream* pByteStream,
    LPCWSTR pURL,
    DWORD dwFlags,
    IPropertyStore*,
    IUnknown** ppCancelCookie,
    IMFAsyncCallback* pCallback,
    IUnknown* pState)
{
    pURL;  //?

#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::BeginCreateObject: url="
       << (pURL ? pURL : L"<NONE>")
       << endl;
#endif

    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

    if (ppCancelCookie)
        *ppCancelCookie = 0;

    if (pByteStream == 0)
        return E_INVALIDARG;

    if (pCallback == 0)
        return E_INVALIDARG;

    if ((dwFlags & MF_RESOLUTION_MEDIASOURCE) == 0)
        return E_INVALIDARG;

    if (m_pResult)  //assume one-at-a-time creation
        return MF_E_UNEXPECTED;

    DWORD dw;
    hr = pByteStream->GetCapabilities(&dw);

    if (SUCCEEDED(hr))
    {
#ifdef _DEBUG
        odbgstream os;
        os << "webmmfsource::BeginCreateObject: caps:";

        if (dw & MFBYTESTREAM_IS_READABLE)
            os << " READABLE";

        if (dw & MFBYTESTREAM_IS_WRITABLE)
            os << " WRITABLE";

        if (dw & MFBYTESTREAM_IS_SEEKABLE)
            os << " SEEKABLE";

        if (dw & MFBYTESTREAM_IS_REMOTE)
            os << " REMOTE";

        if (dw & MFBYTESTREAM_IS_DIRECTORY)
            os << " DIRECTORY";

        if (dw & MFBYTESTREAM_HAS_SLOW_SEEK)
            os << " SLOW_SEEK";

        if (dw & MFBYTESTREAM_IS_PARTIALLY_DOWNLOADED)
            os << " PARTIALLY_DOWNLOADED";

        if (dw & MFBYTESTREAM_SHARE_WRITE)
            os << " SHARE_WRITE";

        os << endl;
#endif

        if (!(dw & MFBYTESTREAM_IS_READABLE))
            return E_FAIL;

        if (!(dw & MFBYTESTREAM_IS_SEEKABLE))
            return MF_E_BYTESTREAM_NOT_SEEKABLE;
    }

    QWORD length;

    hr = pByteStream->GetLength(&length);

    if (FAILED(hr))  //MF_E_BYTESTREAM_UNKNOWN_LENGTH
        return hr;

    if (length == 0)
        return MF_E_INVALID_FILE_FORMAT;  //TODO: use this elsewhere

    WebmMfSource* pSource;

    hr = WebmMfSource::CreateSource(m_pClassFactory, pByteStream, pSource);

    if (FAILED(hr))
        return hr;

    const IMFMediaSourcePtr pUnk(pSource, false);

    IMFAsyncResultPtr pResult;

    hr = MFCreateAsyncResult(pUnk, pCallback, pState, &pResult);

    if (FAILED(hr))
        return hr;

    hr = pSource->BeginLoad(&m_async_load);

    if (FAILED(hr))
        return hr;

    m_pResult = pResult;

    m_bCancel = FALSE;
    m_pByteStream = pByteStream;

    return S_OK;
}

HRESULT WebmMfByteStreamHandler::EndCreateObject(
    IMFAsyncResult* pResult,
    MF_OBJECT_TYPE* pObjectType,
    IUnknown** ppObject)
{
    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::EndCreateObject (begin)" << endl;
#endif

    if (ppObject == 0)
        return E_POINTER;

    IUnknown*& pObject = *ppObject;
    pObject = 0;

    if (pObjectType == 0)
        return E_POINTER;

    MF_OBJECT_TYPE& type = *pObjectType;
    type = MF_OBJECT_INVALID;

    if (pResult == 0)
        return E_INVALIDARG;

    hr = pResult->GetStatus();

    if (FAILED(hr))  //for example, cancelled
        return hr;

    IUnknownPtr pUnk;

    hr = pResult->GetObject(&pUnk);

    if (FAILED(hr))  //shouldn't happen
        return hr;

    assert(pUnk);

    hr = pUnk->QueryInterface(__uuidof(IMFMediaSource), (void**)&pObject);

    if (FAILED(hr))
        return hr;

    type = MF_OBJECT_MEDIASOURCE;

#ifdef _DEBUG
    os << L"WebmMfByteStreamHandler::EndCreateObject (end)" << endl;
#endif

    return S_OK;
}

HRESULT WebmMfByteStreamHandler::CancelObjectCreation(IUnknown*)
{
#ifdef _DEBUG
    wodbgstream os;
    os << L"WebmMfByteStreamHandler::CancelObjectCreation" << endl;
#endif

    Lock lock;

    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;

    if (m_bCancel)
        return S_OK;

    if (!m_pByteStream)
        return S_OK;

    m_bCancel = TRUE;

    return m_pByteStream->Close();
}

HRESULT WebmMfByteStreamHandler::GetMaxNumberOfBytesRequiredForResolution(
    QWORD* pcb)
{
    if (pcb == 0)
        return E_POINTER;

    *pcb = 1024;  //TODO: ask the webm parser for this value

    return S_OK;
}

WebmMfByteStreamHandler::CAsyncLoad::CAsyncLoad(
    WebmMfByteStreamHandler* p) :
    m_pHandler(p)
{
}

WebmMfByteStreamHandler::CAsyncLoad::~CAsyncLoad()
{
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::QueryInterface(
    const IID& iid,
    void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);

    if (iid == __uuidof(IUnknown))
    {
        pUnk = static_cast<IMFAsyncCallback*>(this);  //must be nondelegating
    }
    else if (iid == __uuidof(IMFAsyncCallback))
    {
        pUnk = static_cast<IMFAsyncCallback*>(this);
    }
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }

    pUnk->AddRef();
    return S_OK;
}

ULONG WebmMfByteStreamHandler::CAsyncLoad::AddRef()
{
    return m_pHandler->AddRef();
}

ULONG WebmMfByteStreamHandler::CAsyncLoad::Release()
{
    return m_pHandler->Release();
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::GetParameters(
    DWORD*,
    DWORD*)
{
    return E_NOTIMPL;  //means "assume default behavior"
}

HRESULT WebmMfByteStreamHandler::CAsyncLoad::Invoke(IMFAsyncResult* pResult)
{
    if (pResult == 0)
        return E_INVALIDARG;

    Lock lock;

    HRESULT hr = lock.Seize(m_pHandler);

    if (FAILED(hr))
        return hr;

    return m_pHandler->EndLoad(pResult);
}

HRESULT WebmMfByteStreamHandler::EndLoad(IMFAsyncResult* pLoadResult)
{
    if (!m_pResult)  //should never happen
        return MF_E_UNEXPECTED;

    HRESULT hrLoad;

    if (m_bCancel)
        hrLoad = MF_E_OPERATION_CANCELLED;
    else
        hrLoad = pLoadResult->GetStatus();

    HRESULT hr = m_pResult->SetStatus(hrLoad);
    assert(SUCCEEDED(hr));

    if (FAILED(hrLoad))
    {
        IUnknownPtr pUnk;

        hr = m_pResult->GetObject(&pUnk);

        if (SUCCEEDED(hr))
        {
            const IMFMediaSourcePtr pSource(pUnk);
            assert(pSource);

            hr = pSource->Shutdown();
        }
    }

    hr = MFInvokeCallback(m_pResult);

    m_pByteStream = 0;
    m_pResult = 0;

    return hr;
}

}  //end namespace WebmMfSourceLib
<|endoftext|>


<commit_before><commit_msg>Pull in r182656 from upstream llvm trunk:<commit_after><|endoftext|>


<commit_before>/*
	This file implements the classes defined in Match.h:

	Match, Capability, and Client

	A Match object contains all of the information a startd needs
    about a given match, such as the capability, the client of this
    match, etc.  The capability in the match is just a pointer to a
    Capability object.  The client is also just a pointer to a Client
    object.  The startd maintains two Match objects in the "rip", the
    per-resource information structure.  One for the current match,
    and one for the possibly preempting match that is pending.

	A Capability object contains the capability string, and some
	functions to manipulate and compare against this string.  The
	constructor generates a new capability with the following form:
	<ip:port>#random_integer

	A Client object contains all the info about a given client of a
	startd.  In particular, the client name (a.k.a. "user"), the
	address ("<www.xxx.yyy.zzz:pppp>" formed ip/port), and the
	hostname.

	Written 9/29/97 by Derek Wright <wright@cs.wisc.edu>
*/

#include "startd.h"
static char *_FileName_ = __FILE__;

///////////////////////////////////////////////////////////////////////////
// Match
///////////////////////////////////////////////////////////////////////////

Match::Match()
{
	m_client = new Client;
	m_cap = new Capability;
	m_ad = NULL;
	m_rank = -1;
	m_oldrank = -1;
	m_universe = -1;
	m_agentstream = NULL;
	m_match_tid = -1;
	m_claim_tid = -1;
	m_aliveint = -1;
	m_cluster = -1;
	m_proc = -1;
	m_job_start = -1;
	m_last_pckpt = -1;
}

Match::~Match()
{
		// Cancel any timers associated with this match
	this->cancel_match_timer();
	this->cancel_claim_timer();

		// Free up memory that's been allocated
	if( m_ad ) {
		delete( m_ad );
	}
	delete( m_cap );
	if( m_client ) {
		delete( m_client );
	}
	if( m_agentstream ) {
		delete( m_agentstream );
	}

}

void
Match::vacate()
{
	assert( m_cap );
		// warn the client of this match that it's being vacated
	if( m_client && m_client->addr() ) {
		m_client->vacate( m_cap->capab() );
	}
}

int
Match::send_accountant( int cmd )
{
	if( !accountant_host ) {
		return FALSE;
	}
	ReliSock sock;
	sock.timeout( 30 );
	if( ! sock.connect( accountant_host, ACCOUNTANT_PORT ) ) {
		dprintf(D_ALWAYS, "Couldn't connect to accountant\n");
		return FALSE;
	}
	if( !sock.put( cmd ) ) {
		dprintf(D_ALWAYS, "Can't send command (%d) to accountant\n", cmd );
		return FALSE;
	}
	if( !sock.put( m_cap->capab() ) ) {
		dprintf(D_ALWAYS, "Can't send capability to accountant\n");
		return FALSE;
	}
	if( !sock.eom() ) {
		dprintf(D_ALWAYS, "Can't send EOM to accountant\n");
		return FALSE;
	}
	sock.close();
	return TRUE;
}

void
Match::update( ClassAd* ad )
{
	char line[256];
	char* tmp;

	sprintf( line, "%s = %f", ATTR_CURRENT_RANK, m_rank );
	ad->Insert( line );

	if( m_client ) {
		tmp = m_client->name();
		if( tmp ) {
			sprintf( line, "%s=\"%s\"", ATTR_REMOTE_USER, tmp );
			ad->Insert( line );
		}
		tmp = m_client->host();
		if( tmp ) {
			sprintf( line, "%s=\"%s\"", ATTR_CLIENT_MACHINE, tmp );
			ad->Insert( line );
		}
	}

	if( (m_cluster > 0) && (m_proc >= 0) ) {
		sprintf( line, "%s=\"%d.%d\"", ATTR_JOB_ID, m_cluster, m_proc );
		ad->Insert( line );
	}

	if( m_job_start > 0 ) {
		sprintf(line, "%s=%d", ATTR_JOB_START, m_job_start );
		ad->Insert( line );
	}

	if( m_last_pckpt > 0 ) {
		sprintf(line, "%s=%d", ATTR_LAST_PERIODIC_CHECKPOINT, m_last_pckpt );
		ad->Insert( line );
	}

	if( m_universe > 0 ) {
		sprintf( line, "%s=%d", ATTR_JOB_UNIVERSE, m_universe );
		ad->Insert( line );
	}
}

void
Match::refuse_agent()
{
	if( !m_agentstream ) return;
	dprintf( D_ALWAYS, "Refusing request from schedd agent.\n" );
	m_agentstream->encode();
	m_agentstream->put(NOT_OK);
	m_agentstream->end_of_message();
}

void
Match::start_match_timer()
{
	m_match_tid =
		daemonCore->Register_Timer( match_timeout, 0,
								   (TimerHandlercpp)match_timed_out,
								   "match_timed_out", this );
	if( m_match_tid == -1 ) {
		EXCEPT( "Couldn't register timer (out of memory)." );
	}
	dprintf( D_FULLDEBUG, "Started match timer for %d seconds.\n",
			 match_timeout );
}

void
Match::cancel_match_timer()
{
	if( m_match_tid != -1 ) {
		daemonCore->Cancel_Timer( m_match_tid );
		m_match_tid = -1;
		dprintf( D_FULLDEBUG, "Canceled match timer.\n" );
	}
}

int
Match::match_timed_out()
{
	Resource* rip = resmgr->get_by_any_cap( capab() );
	if( !rip ) {
		EXCEPT( "Can't find resource of expired match." );
	}

	if( rip->r_cur->cap()->matches( capab() ) ) {
		if( rip->state() != matched_state ) {
			EXCEPT( "Match timed out but not in matched state." );
		}
		delete rip->r_cur;
		rip->r_cur = new Match;
		rip->change_state( owner_state );
	} else {
			// The match that timed out was the preempting match.
		assert( rip->r_pre->cap()->matches( capab() ) );
			// We need to generate a new preempting match object,
			// restore our reqexp, and update the CM.
		delete rip->r_pre;
		rip->r_pre = new Match;
		rip->r_reqexp->pub();
		rip->update();
	}
	return TRUE;
}

void
Match::start_claim_timer()
{
	if( m_aliveint < 0 ) {
		dprintf( D_ALWAYS,
				 "Warning: starting claim timer before alive interval set.\n" );
		m_aliveint = 300;
	}
	m_claim_tid =
		daemonCore->Register_Timer( (3 * m_aliveint), 0,
									(TimerHandlercpp)claim_timed_out,
									"claim_timed_out", this );
	if( m_claim_tid == -1 ) {
		EXCEPT( "Couldn't register timer (out of memory)." );
	}
	dprintf( D_FULLDEBUG, "Started claim timer w/ %d second alive interval.\n",
			 m_aliveint );
}

void
Match::cancel_claim_timer()
{
	if( m_claim_tid != -1 ) {
		daemonCore->Cancel_Timer( m_claim_tid );
		m_claim_tid = -1;
		dprintf( D_FULLDEBUG, "Canceled claim timer.\n" );
	}
}

int
Match::claim_timed_out()
{
	Resource* rip = resmgr->get_by_cur_cap( capab() );
	if( !rip ) {
		EXCEPT( "Can't find resource of expired claim." );
	}
		// Note that this claim timed out so we don't try to send a
		// command to our client.
	if( m_client ) {
		delete m_client;
		m_client = NULL;
	}
		// Release the claim.
	rip->release_claim();
	return TRUE;
}

void
Match::alive()
{
		// Process a keep alive command
	daemonCore->Reset_Timer( m_claim_tid, (3 * m_aliveint), 0 );
}

// Set our ad to the given pointer
void
Match::setad(ClassAd *ad)
{
	if( m_ad ) {
		delete( m_ad );
	}
	m_ad = ad;
}

void
Match::deletead(void)
{
	if( m_ad ) {
		delete( m_ad );
		m_ad = NULL;
	}
}

void
Match::setagentstream(Stream* stream)
{
	if( m_agentstream ) {
		delete( m_agentstream );
	}
	m_agentstream = stream;
}

///////////////////////////////////////////////////////////////////////////
// Client
///////////////////////////////////////////////////////////////////////////

Client::Client()
{
	c_name = NULL;
	c_addr = NULL;
	c_host = NULL;
}

Client::~Client()
{
	free( c_name );
	free( c_addr );
	free( c_host );
}

void
Client::setname(char* name)
{
	if( c_name ) {
		free( c_name);
	}
	if( name ) {
		c_name = strdup( name );
	} else {
		c_name = NULL;
	}
}

void
Client::setaddr(char* addr)
{
	if( c_addr ) {
		free( c_addr);
	}
	if( addr ) {
		c_addr = strdup( addr );
	} else {
		c_addr = NULL;
	}
}

void
Client::sethost(char* host)
{
	if( c_host ) {
		free( c_host);
	}
	if( host ) {
		c_host = strdup( host );
	} else {
		c_host = NULL;
	}
}

void
Client::vacate(char* cap)
{
	ReliSock sock;
	sock.timeout( 20 );

	if( ! (c_addr || c_host || c_name ) ) {
			// Client not really set, nothing to do.
		return;
	}

	dprintf(D_FULLDEBUG, "Entered vacate_client %s %s...\n", c_addr, c_host);

	if(	! sock.connect( c_addr, 0 ) ) {
		dprintf(D_ALWAYS, "Can't connect to schedd (%s)\n", c_addr);
	} else {
		if( !sock.put( RELEASE_CLAIM ) ) {
			dprintf(D_ALWAYS, "Can't send RELEASE_CLAIM command to client\n");
		} else if( !sock.put( cap ) ) {
			dprintf(D_ALWAYS, "Can't send capability to client\n");
		} else if( !sock.eom() ) {
			dprintf(D_ALWAYS, "Can't send EOM to client\n");
		}
		sock.close();
	}
}

///////////////////////////////////////////////////////////////////////////
// Capability
///////////////////////////////////////////////////////////////////////////

char*
new_capability_string()
{
	char cap[128];
	char tmp[128];
	char randbuf[12];
	randbuf[0] = '\0';
	int i, len;

		// Create a really mangled 10 digit random number: The first 6
		// digits are generated as follows: for the ith digit, pull
		// the ith digit off a new random int.  So our 1st slot comes
		// from the 1st digit of 1 random int, the 2nd from the 2nd
		// digit of a 2nd random it, etc...  If we're trying to get a
		// digit from a number that's too small to have that many, we
		// just use the last digit.  The last 4 digits of our number
		// come from the first 4 digits of the current time multiplied
		// by a final random int.  That should keep 'em guessing. :)
		// -Derek Wright 1/8/98
	for( i=0; i<6; i++ ) {
		sprintf( tmp, "%d", get_random_int() );
		len = strlen(tmp);
		if( i < len ) {
			tmp[i+1] = '\0';
			strcat( randbuf, tmp+i );
		} else {
			strcat( randbuf, tmp+(len-1) );
		}
	}
	sprintf( tmp, "%f", (double)((float)time(NULL) * (float)get_random_int()) );
	tmp[4]='\0';
	strcat( randbuf, tmp );

		// Capability string is "<ip:port>#random_number"
	strcpy( cap, daemonCore->InfoCommandSinfulString() );
	strcat( cap, "#" );
	strcat( cap, randbuf );
	return strdup( cap );
}

Capability::Capability()
{
	c_capab = new_capability_string();
}

Capability::~Capability()
{
	free( c_capab );
}

int
Capability::matches(char* capab)
{
	return( strcmp(capab, c_capab) == 0 );
}

<commit_msg>Don't put ATTR_JOB_UNIVERSE into the resource classad.  We don't use it, and if the users want it in the classad, they can put it into startd_job_exprs.<commit_after>/*
	This file implements the classes defined in Match.h:

	Match, Capability, and Client

	A Match object contains all of the information a startd needs
    about a given match, such as the capability, the client of this
    match, etc.  The capability in the match is just a pointer to a
    Capability object.  The client is also just a pointer to a Client
    object.  The startd maintains two Match objects in the "rip", the
    per-resource information structure.  One for the current match,
    and one for the possibly preempting match that is pending.

	A Capability object contains the capability string, and some
	functions to manipulate and compare against this string.  The
	constructor generates a new capability with the following form:
	<ip:port>#random_integer

	A Client object contains all the info about a given client of a
	startd.  In particular, the client name (a.k.a. "user"), the
	address ("<www.xxx.yyy.zzz:pppp>" formed ip/port), and the
	hostname.

	Written 9/29/97 by Derek Wright <wright@cs.wisc.edu>
*/

#include "startd.h"
static char *_FileName_ = __FILE__;

///////////////////////////////////////////////////////////////////////////
// Match
///////////////////////////////////////////////////////////////////////////

Match::Match()
{
	m_client = new Client;
	m_cap = new Capability;
	m_ad = NULL;
	m_rank = -1;
	m_oldrank = -1;
	m_universe = -1;
	m_agentstream = NULL;
	m_match_tid = -1;
	m_claim_tid = -1;
	m_aliveint = -1;
	m_cluster = -1;
	m_proc = -1;
	m_job_start = -1;
	m_last_pckpt = -1;
}

Match::~Match()
{
		// Cancel any timers associated with this match
	this->cancel_match_timer();
	this->cancel_claim_timer();

		// Free up memory that's been allocated
	if( m_ad ) {
		delete( m_ad );
	}
	delete( m_cap );
	if( m_client ) {
		delete( m_client );
	}
	if( m_agentstream ) {
		delete( m_agentstream );
	}

}

void
Match::vacate()
{
	assert( m_cap );
		// warn the client of this match that it's being vacated
	if( m_client && m_client->addr() ) {
		m_client->vacate( m_cap->capab() );
	}
}

int
Match::send_accountant( int cmd )
{
	if( !accountant_host ) {
		return FALSE;
	}
	ReliSock sock;
	sock.timeout( 30 );
	if( ! sock.connect( accountant_host, ACCOUNTANT_PORT ) ) {
		dprintf(D_ALWAYS, "Couldn't connect to accountant\n");
		return FALSE;
	}
	if( !sock.put( cmd ) ) {
		dprintf(D_ALWAYS, "Can't send command (%d) to accountant\n", cmd );
		return FALSE;
	}
	if( !sock.put( m_cap->capab() ) ) {
		dprintf(D_ALWAYS, "Can't send capability to accountant\n");
		return FALSE;
	}
	if( !sock.eom() ) {
		dprintf(D_ALWAYS, "Can't send EOM to accountant\n");
		return FALSE;
	}
	sock.close();
	return TRUE;
}

void
Match::update( ClassAd* ad )
{
	char line[256];
	char* tmp;

	sprintf( line, "%s = %f", ATTR_CURRENT_RANK, m_rank );
	ad->Insert( line );

	if( m_client ) {
		tmp = m_client->name();
		if( tmp ) {
			sprintf( line, "%s=\"%s\"", ATTR_REMOTE_USER, tmp );
			ad->Insert( line );
		}
		tmp = m_client->host();
		if( tmp ) {
			sprintf( line, "%s=\"%s\"", ATTR_CLIENT_MACHINE, tmp );
			ad->Insert( line );
		}
	}

	if( (m_cluster > 0) && (m_proc >= 0) ) {
		sprintf( line, "%s=\"%d.%d\"", ATTR_JOB_ID, m_cluster, m_proc );
		ad->Insert( line );
	}

	if( m_job_start > 0 ) {
		sprintf(line, "%s=%d", ATTR_JOB_START, m_job_start );
		ad->Insert( line );
	}

	if( m_last_pckpt > 0 ) {
		sprintf(line, "%s=%d", ATTR_LAST_PERIODIC_CHECKPOINT, m_last_pckpt );
		ad->Insert( line );
	}
}

void
Match::refuse_agent()
{
	if( !m_agentstream ) return;
	dprintf( D_ALWAYS, "Refusing request from schedd agent.\n" );
	m_agentstream->encode();
	m_agentstream->put(NOT_OK);
	m_agentstream->end_of_message();
}

void
Match::start_match_timer()
{
	m_match_tid =
		daemonCore->Register_Timer( match_timeout, 0,
								   (TimerHandlercpp)match_timed_out,
								   "match_timed_out", this );
	if( m_match_tid == -1 ) {
		EXCEPT( "Couldn't register timer (out of memory)." );
	}
	dprintf( D_FULLDEBUG, "Started match timer for %d seconds.\n",
			 match_timeout );
}

void
Match::cancel_match_timer()
{
	if( m_match_tid != -1 ) {
		daemonCore->Cancel_Timer( m_match_tid );
		m_match_tid = -1;
		dprintf( D_FULLDEBUG, "Canceled match timer.\n" );
	}
}

int
Match::match_timed_out()
{
	Resource* rip = resmgr->get_by_any_cap( capab() );
	if( !rip ) {
		EXCEPT( "Can't find resource of expired match." );
	}

	if( rip->r_cur->cap()->matches( capab() ) ) {
		if( rip->state() != matched_state ) {
			EXCEPT( "Match timed out but not in matched state." );
		}
		delete rip->r_cur;
		rip->r_cur = new Match;
		rip->change_state( owner_state );
	} else {
			// The match that timed out was the preempting match.
		assert( rip->r_pre->cap()->matches( capab() ) );
			// We need to generate a new preempting match object,
			// restore our reqexp, and update the CM.
		delete rip->r_pre;
		rip->r_pre = new Match;
		rip->r_reqexp->pub();
		rip->update();
	}
	return TRUE;
}

void
Match::start_claim_timer()
{
	if( m_aliveint < 0 ) {
		dprintf( D_ALWAYS,
				 "Warning: starting claim timer before alive interval set.\n" );
		m_aliveint = 300;
	}
	m_claim_tid =
		daemonCore->Register_Timer( (3 * m_aliveint), 0,
									(TimerHandlercpp)claim_timed_out,
									"claim_timed_out", this );
	if( m_claim_tid == -1 ) {
		EXCEPT( "Couldn't register timer (out of memory)." );
	}
	dprintf( D_FULLDEBUG, "Started claim timer w/ %d second alive interval.\n",
			 m_aliveint );
}

void
Match::cancel_claim_timer()
{
	if( m_claim_tid != -1 ) {
		daemonCore->Cancel_Timer( m_claim_tid );
		m_claim_tid = -1;
		dprintf( D_FULLDEBUG, "Canceled claim timer.\n" );
	}
}

int
Match::claim_timed_out()
{
	Resource* rip = resmgr->get_by_cur_cap( capab() );
	if( !rip ) {
		EXCEPT( "Can't find resource of expired claim." );
	}
		// Note that this claim timed out so we don't try to send a
		// command to our client.
	if( m_client ) {
		delete m_client;
		m_client = NULL;
	}
		// Release the claim.
	rip->release_claim();
	return TRUE;
}

void
Match::alive()
{
		// Process a keep alive command
	daemonCore->Reset_Timer( m_claim_tid, (3 * m_aliveint), 0 );
}

// Set our ad to the given pointer
void
Match::setad(ClassAd *ad)
{
	if( m_ad ) {
		delete( m_ad );
	}
	m_ad = ad;
}

void
Match::deletead(void)
{
	if( m_ad ) {
		delete( m_ad );
		m_ad = NULL;
	}
}

void
Match::setagentstream(Stream* stream)
{
	if( m_agentstream ) {
		delete( m_agentstream );
	}
	m_agentstream = stream;
}

///////////////////////////////////////////////////////////////////////////
// Client
///////////////////////////////////////////////////////////////////////////

Client::Client()
{
	c_name = NULL;
	c_addr = NULL;
	c_host = NULL;
}

Client::~Client()
{
	free( c_name );
	free( c_addr );
	free( c_host );
}

void
Client::setname(char* name)
{
	if( c_name ) {
		free( c_name);
	}
	if( name ) {
		c_name = strdup( name );
	} else {
		c_name = NULL;
	}
}

void
Client::setaddr(char* addr)
{
	if( c_addr ) {
		free( c_addr);
	}
	if( addr ) {
		c_addr = strdup( addr );
	} else {
		c_addr = NULL;
	}
}

void
Client::sethost(char* host)
{
	if( c_host ) {
		free( c_host);
	}
	if( host ) {
		c_host = strdup( host );
	} else {
		c_host = NULL;
	}
}

void
Client::vacate(char* cap)
{
	ReliSock sock;
	sock.timeout( 20 );

	if( ! (c_addr || c_host || c_name ) ) {
			// Client not really set, nothing to do.
		return;
	}

	dprintf(D_FULLDEBUG, "Entered vacate_client %s %s...\n", c_addr, c_host);

	if(	! sock.connect( c_addr, 0 ) ) {
		dprintf(D_ALWAYS, "Can't connect to schedd (%s)\n", c_addr);
	} else {
		if( !sock.put( RELEASE_CLAIM ) ) {
			dprintf(D_ALWAYS, "Can't send RELEASE_CLAIM command to client\n");
		} else if( !sock.put( cap ) ) {
			dprintf(D_ALWAYS, "Can't send capability to client\n");
		} else if( !sock.eom() ) {
			dprintf(D_ALWAYS, "Can't send EOM to client\n");
		}
		sock.close();
	}
}

///////////////////////////////////////////////////////////////////////////
// Capability
///////////////////////////////////////////////////////////////////////////

char*
new_capability_string()
{
	char cap[128];
	char tmp[128];
	char randbuf[12];
	randbuf[0] = '\0';
	int i, len;

		// Create a really mangled 10 digit random number: The first 6
		// digits are generated as follows: for the ith digit, pull
		// the ith digit off a new random int.  So our 1st slot comes
		// from the 1st digit of 1 random int, the 2nd from the 2nd
		// digit of a 2nd random it, etc...  If we're trying to get a
		// digit from a number that's too small to have that many, we
		// just use the last digit.  The last 4 digits of our number
		// come from the first 4 digits of the current time multiplied
		// by a final random int.  That should keep 'em guessing. :)
		// -Derek Wright 1/8/98
	for( i=0; i<6; i++ ) {
		sprintf( tmp, "%d", get_random_int() );
		len = strlen(tmp);
		if( i < len ) {
			tmp[i+1] = '\0';
			strcat( randbuf, tmp+i );
		} else {
			strcat( randbuf, tmp+(len-1) );
		}
	}
	sprintf( tmp, "%f", (double)((float)time(NULL) * (float)get_random_int()) );
	tmp[4]='\0';
	strcat( randbuf, tmp );

		// Capability string is "<ip:port>#random_number"
	strcpy( cap, daemonCore->InfoCommandSinfulString() );
	strcat( cap, "#" );
	strcat( cap, randbuf );
	return strdup( cap );
}

Capability::Capability()
{
	c_capab = new_capability_string();
}

Capability::~Capability()
{
	free( c_capab );
}

int
Capability::matches(char* capab)
{
	return( strcmp(capab, c_capab) == 0 );
}

<|endoftext|>


<commit_before>/*
 * Copyright(c) Sophist Solutions, Inc. 1990-2011.  All rights reserved
 */
#include	"../StroikaPreComp.h"

#include	<algorithm>
#include	<cstdlib>
#include	<fstream>
#include	<sstream>

#if		qPlatform_POSIX
	#include	<sys/types.h>
	#include	<unistd.h>
	#include	<signal.h>
#endif

#include	"../../Foundation/Characters/Format.h"
#include	"../../Foundation/Containers/Common.h"
#include	"../../Foundation/Debug/Assertions.h"
#include	"../../Foundation/Debug/Trace.h"
#include	"../../Foundation/Execution/CommandLine.h"
#include	"../../Foundation/Execution/Exceptions.h"
#include	"../../Foundation/Execution/Module.h"
#include	"../../Foundation/Execution/ThreadAbortException.h"
#include	"../../Foundation/IO/FileUtils.h"
#include	"../../Foundation/Memory/SmallStackBuffer.h"

#include	"Main.h"

using	namespace	Stroika::Foundation;
using	namespace	Stroika::Foundation::Containers;
using	namespace	Stroika::Foundation::Memory;

using	namespace	Stroika::Frameworks;
using	namespace	Stroika::Frameworks::Service;

/*
 ********************************************************************************
 ****************************** Service::Main::IRep *****************************
 ********************************************************************************
 */
Main::IRep::IRep ()
	: fStopping_ (false)
	, fMustReReadConfig (false)
{
}

Main::IRep::~IRep ()
{
}

void	Main::IRep::OnStartRequest ()
{
	// This procedure ends when the entire service process ends...
	Debug::TraceContextBumper traceCtx (TSTR ("Stroika::Frameworks::Service::Main::IRep::OnStartRequest"));
	MainLoop ();
}

void	Main::IRep::OnStopRequest ()
{
	// default to using thread stuff to send us a signal to abort...
	fStopping_ = true;
}

void	Main::IRep::OnReReadConfigurationRequest ()
{
}

String	Main::IRep::GetServiceStatusMessage () const
{
	return String ();
}

#if		qPlatform_POSIX
String	Main::IRep::GetPIDFileName () const
{
	return L"/tmp/" + GetServiceDescription ().fName + L".pid";
}
#endif

#if		qPlatform_POSIX
void	Main::IRep::SignalHandler (int signum)
{
	// VERY PRIMITIVE IMPL FOR NOW -- LGP 2011-09-24
	switch (signum) {
	case	SIGTERM:
		fStopping_ = true;
		break;
	#if		qCompilerAndStdLib_Supports_constexpr
		case	kSIG_ReReadConfiguration:
	#else
		case	SIGHUP:
	#endif
			fMustReReadConfig = true;
			break;

	}
}
#endif

/*
 ********************************************************************************
 ************************************ Service::Main *****************************
 ********************************************************************************
 */

const	wchar_t	Service::Main::CommandNames::kRunAsService[]		=	L"Run-As-Service";
const	wchar_t	Service::Main::CommandNames::kStart[]				=	L"Start";
const	wchar_t	Service::Main::CommandNames::kStop[]				=	L"Stop";
const	wchar_t	Service::Main::CommandNames::kKill[]				=	L"Kill";
const	wchar_t	Service::Main::CommandNames::kRestart[]				=	L"Restart";
const	wchar_t	Service::Main::CommandNames::kReloadConfiguration[]	=	L"Reload-Configuration";
const	wchar_t	Service::Main::CommandNames::kPause[]				=	L"Pause";
const	wchar_t	Service::Main::CommandNames::kContinue[]			=	L"Continue";

Memory::SharedPtr<Main::IRep>	Main::_sRep;

Main::Main (Memory::SharedPtr<IRep> rep)
{
	Ensure  (_sRep.IsNull ());
	_sRep = rep;
	#if		qPlatform_POSIX
		SetupSignalHanlders_ ();
	#endif
}

#if		qPlatform_POSIX
void	Main::SetupSignalHanlders_ ()
{
	signal (SIGTERM, SignalHandler);
	signal (kSIG_ReReadConfiguration, SignalHandler);
}
#endif

Main::State	Main::GetState () const
{
	#if		qPlatform_POSIX
	if (GetServicePID () != 0) {
		return eRunning;
	}
	#endif
	return eStopped;	// otherwise (esp on other platforms where not implemtned) must  be stopped
}

#if		qPlatform_POSIX
pid_t	Main::GetServicePID () const
{
	ifstream	in (_sRep->GetPIDFileName ().AsTString ().c_str ());
	if (in) {
		pid_t	n = 0;
		in >> n;
		return n;
	}
	return 0;
}
#endif

String		Main::GetServiceStatusMessage () const
{
	ServiceDescription	svd	=	GetServiceDescription ();
	wstringstream	tmp;
	tmp << L"Service '" << svd.fName.As<wstring> () << "'" << endl;
	switch (this->GetState ()) {
		case	eStopped:
			tmp << L"    State:    STOPPED" << endl;
			break;
		case	eRunning:
			tmp << L"    State:    Running" << endl;
			#if		qPlatform_POSIX
				tmp << L"    PID:    " << GetServicePID () << endl;
			#endif
			break;
		case	ePaused:
			tmp << L"    State:    PAUSED" << endl;
			#if		qPlatform_POSIX
				tmp << L"    PID:    " << GetServicePID () << endl;
			#endif
			break;
		default:
			AssertNotReached ();
	}
	return tmp.str ();
}

void	Main::RunAsService ()
{
	// VERY PRIMITIVE IMPL - WE NEED LOCKING on tmpfile stuff - add good tempfile supprot to fileuitls or use existing...

	try {
#if		qPlatform_POSIX
		ofstream	out (_sRep->GetPIDFileName ().AsTString ().c_str ());
		out << getpid () << endl;
#endif
		_sRep->OnStartRequest ();
	}
	catch (const Execution::ThreadAbortException& /*threadAbort*/) {
#if		qPlatform_POSIX
		unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
		// ignore this - just means service ended normally
	}
	catch (...) {
		DbgTrace (TSTR ("Unexpected exception ended running service"));
#if		qPlatform_POSIX
		unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
		throw;
	}
}

void	Main::Start ()
{
	// Check not already runnig, (someday) and then for and exec the

#if		qPlatform_POSIX
	// REALLY should use GETSTATE - and return state based on if PID file exsits...
	if (GetServicePID ()  != 0) {
		Execution::DoThrow (Execution::StringException (L"Cannot Start service because its already running"));
	}
#endif
	Characters::TString	thisEXEPath	=	Execution::GetEXEPath ();
#if		qPlatform_POSIX
	pid_t	pid	=	fork ();
	if (pid == 0) {
		// Child - exec
		int	r	=	execl (thisEXEPath.c_str (), thisEXEPath.c_str (), String (CommandNames::kRunAsService).AsTString ().c_str (), nullptr);
		exit (-1);
	}
	else if (pid < 0) {
		// failed to fork - serious error
		Execution::errno_ErrorException::DoThrow (errno);
	}
	else {
		// parent - in this case - no reason to wait - our work is done... Future versions might wait to
		// see if the 'pidfile' got created....
		//		--LGP 2011-09-23
	}
#endif
}

void	Main::Stop ()
{
	// Send signal to server to stop
#if		qPlatform_POSIX
	kill (GetServicePID (), SIGTERM);
#endif
}

void	Main::Kill ()
{
	Stop ();
	// Send signal to server to stop
#if		qPlatform_POSIX
	kill (GetServicePID (), SIGKILL);
	// REALY should WAIT for server to stop and only do this it fails -
	unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
}

void	Main::Restart ()
{
	IgnoreExceptionsForCall (Stop ());
#if		qPlatform_POSIX
	// REALY should WAIT for server to stop and only do this it fails -
	unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
	Start ();
}

void	Main::ReReadConfiguration ()
{
	// SEND APPROPRIATE SIGNAL
#if		qPlatform_POSIX
	pid_t	pid	=	GetServicePID ();
	Assert (pid != 0);	// maybe throw if non-zero???
	kill (GetServicePID (), kSIG_ReReadConfiguration);
#endif
}

void	Main::Pause ()
{
	AssertNotImplemented ();
}

void	Main::Continue ()
{
	AssertNotImplemented ();
}

Main::ServiceDescription	Main::GetServiceDescription () const
{
	return _sRep->GetServiceDescription ();
}

#if		qPlatform_POSIX
void	Main::SignalHandler (int signum)
{
	_sRep->SignalHandler (signum);
}
#endif

bool	Main::_HandleStandardCommandLineArgument (const String& arg)
{
	if (Execution::MatchesCommandLineArgument (arg, CommandNames::kRunAsService)) {
		RunAsService ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kStart)) {
		Start ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kStop)) {
		Stop ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kKill)) {
		Kill ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kRestart)) {
		Restart ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kReloadConfiguration)) {
		ReReadConfiguration ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kPause)) {
		Pause ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kContinue)) {
		Continue ();
		return true;
	}
	/// MANY more neeeded, and fix to use named constants...
	return false;
}

<commit_msg>neaten output of Main::GetServiceStatusMessage ()<commit_after>/*
 * Copyright(c) Sophist Solutions, Inc. 1990-2011.  All rights reserved
 */
#include	"../StroikaPreComp.h"

#include	<algorithm>
#include	<cstdlib>
#include	<fstream>
#include	<sstream>

#if		qPlatform_POSIX
	#include	<sys/types.h>
	#include	<unistd.h>
	#include	<signal.h>
#endif

#include	"../../Foundation/Characters/Format.h"
#include	"../../Foundation/Containers/Common.h"
#include	"../../Foundation/Debug/Assertions.h"
#include	"../../Foundation/Debug/Trace.h"
#include	"../../Foundation/Execution/CommandLine.h"
#include	"../../Foundation/Execution/Exceptions.h"
#include	"../../Foundation/Execution/Module.h"
#include	"../../Foundation/Execution/ThreadAbortException.h"
#include	"../../Foundation/IO/FileUtils.h"
#include	"../../Foundation/Memory/SmallStackBuffer.h"

#include	"Main.h"

using	namespace	Stroika::Foundation;
using	namespace	Stroika::Foundation::Containers;
using	namespace	Stroika::Foundation::Memory;

using	namespace	Stroika::Frameworks;
using	namespace	Stroika::Frameworks::Service;

/*
 ********************************************************************************
 ****************************** Service::Main::IRep *****************************
 ********************************************************************************
 */
Main::IRep::IRep ()
	: fStopping_ (false)
	, fMustReReadConfig (false)
{
}

Main::IRep::~IRep ()
{
}

void	Main::IRep::OnStartRequest ()
{
	// This procedure ends when the entire service process ends...
	Debug::TraceContextBumper traceCtx (TSTR ("Stroika::Frameworks::Service::Main::IRep::OnStartRequest"));
	MainLoop ();
}

void	Main::IRep::OnStopRequest ()
{
	// default to using thread stuff to send us a signal to abort...
	fStopping_ = true;
}

void	Main::IRep::OnReReadConfigurationRequest ()
{
}

String	Main::IRep::GetServiceStatusMessage () const
{
	return String ();
}

#if		qPlatform_POSIX
String	Main::IRep::GetPIDFileName () const
{
	return L"/tmp/" + GetServiceDescription ().fName + L".pid";
}
#endif

#if		qPlatform_POSIX
void	Main::IRep::SignalHandler (int signum)
{
	// VERY PRIMITIVE IMPL FOR NOW -- LGP 2011-09-24
	switch (signum) {
	case	SIGTERM:
		fStopping_ = true;
		break;
	#if		qCompilerAndStdLib_Supports_constexpr
		case	kSIG_ReReadConfiguration:
	#else
		case	SIGHUP:
	#endif
			fMustReReadConfig = true;
			break;

	}
}
#endif

/*
 ********************************************************************************
 ************************************ Service::Main *****************************
 ********************************************************************************
 */

const	wchar_t	Service::Main::CommandNames::kRunAsService[]		=	L"Run-As-Service";
const	wchar_t	Service::Main::CommandNames::kStart[]				=	L"Start";
const	wchar_t	Service::Main::CommandNames::kStop[]				=	L"Stop";
const	wchar_t	Service::Main::CommandNames::kKill[]				=	L"Kill";
const	wchar_t	Service::Main::CommandNames::kRestart[]				=	L"Restart";
const	wchar_t	Service::Main::CommandNames::kReloadConfiguration[]	=	L"Reload-Configuration";
const	wchar_t	Service::Main::CommandNames::kPause[]				=	L"Pause";
const	wchar_t	Service::Main::CommandNames::kContinue[]			=	L"Continue";

Memory::SharedPtr<Main::IRep>	Main::_sRep;

Main::Main (Memory::SharedPtr<IRep> rep)
{
	Ensure  (_sRep.IsNull ());
	_sRep = rep;
	#if		qPlatform_POSIX
		SetupSignalHanlders_ ();
	#endif
}

#if		qPlatform_POSIX
void	Main::SetupSignalHanlders_ ()
{
	signal (SIGTERM, SignalHandler);
	signal (kSIG_ReReadConfiguration, SignalHandler);
}
#endif

Main::State	Main::GetState () const
{
	#if		qPlatform_POSIX
	if (GetServicePID () != 0) {
		return eRunning;
	}
	#endif
	return eStopped;	// otherwise (esp on other platforms where not implemtned) must  be stopped
}

#if		qPlatform_POSIX
pid_t	Main::GetServicePID () const
{
	ifstream	in (_sRep->GetPIDFileName ().AsTString ().c_str ());
	if (in) {
		pid_t	n = 0;
		in >> n;
		return n;
	}
	return 0;
}
#endif

String		Main::GetServiceStatusMessage () const
{
	const	wchar_t	kTAB[]	=	L"    ";	// use spaces instead of tab so formatting independent of tabstop settings
	ServiceDescription	svd	=	GetServiceDescription ();
	wstringstream	tmp;
	tmp << L"Service '" << svd.fName.As<wstring> () << "'" << endl;
	switch (this->GetState ()) {
		case	eStopped:
			tmp << kTAB << L"State:  " << kTAB << "STOPPED" << endl;
			break;
		case	eRunning:
			tmp << kTAB << L"State:  " << kTAB << "Running" << endl;
			#if		qPlatform_POSIX
				tmp << kTAB << L"PID:    " << kTAB << GetServicePID () << endl;
			#endif
			break;
		case	ePaused:
			tmp << kTAB << L"State:  " << kTAB << "PAUSED" << endl;
			#if		qPlatform_POSIX
				tmp << kTAB << L"PID:    " << kTAB<< GetServicePID () << endl;
			#endif
			break;
		default:
			AssertNotReached ();
	}
	return tmp.str ();
}

void	Main::RunAsService ()
{
	// VERY PRIMITIVE IMPL - WE NEED LOCKING on tmpfile stuff - add good tempfile supprot to fileuitls or use existing...

	try {
#if		qPlatform_POSIX
		ofstream	out (_sRep->GetPIDFileName ().AsTString ().c_str ());
		out << getpid () << endl;
#endif
		_sRep->OnStartRequest ();
	}
	catch (const Execution::ThreadAbortException& /*threadAbort*/) {
#if		qPlatform_POSIX
		unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
		// ignore this - just means service ended normally
	}
	catch (...) {
		DbgTrace (TSTR ("Unexpected exception ended running service"));
#if		qPlatform_POSIX
		unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
		throw;
	}
}

void	Main::Start ()
{
	// Check not already runnig, (someday) and then for and exec the

#if		qPlatform_POSIX
	// REALLY should use GETSTATE - and return state based on if PID file exsits...
	if (GetServicePID ()  != 0) {
		Execution::DoThrow (Execution::StringException (L"Cannot Start service because its already running"));
	}
#endif
	Characters::TString	thisEXEPath	=	Execution::GetEXEPath ();
#if		qPlatform_POSIX
	pid_t	pid	=	fork ();
	if (pid == 0) {
		// Child - exec
		int	r	=	execl (thisEXEPath.c_str (), thisEXEPath.c_str (), String (CommandNames::kRunAsService).AsTString ().c_str (), nullptr);
		exit (-1);
	}
	else if (pid < 0) {
		// failed to fork - serious error
		Execution::errno_ErrorException::DoThrow (errno);
	}
	else {
		// parent - in this case - no reason to wait - our work is done... Future versions might wait to
		// see if the 'pidfile' got created....
		//		--LGP 2011-09-23
	}
#endif
}

void	Main::Stop ()
{
	// Send signal to server to stop
#if		qPlatform_POSIX
	kill (GetServicePID (), SIGTERM);
#endif
}

void	Main::Kill ()
{
	Stop ();
	// Send signal to server to stop
#if		qPlatform_POSIX
	kill (GetServicePID (), SIGKILL);
	// REALY should WAIT for server to stop and only do this it fails -
	unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
}

void	Main::Restart ()
{
	IgnoreExceptionsForCall (Stop ());
#if		qPlatform_POSIX
	// REALY should WAIT for server to stop and only do this it fails -
	unlink (_sRep->GetPIDFileName ().AsTString ().c_str ());
#endif
	Start ();
}

void	Main::ReReadConfiguration ()
{
	// SEND APPROPRIATE SIGNAL
#if		qPlatform_POSIX
	pid_t	pid	=	GetServicePID ();
	Assert (pid != 0);	// maybe throw if non-zero???
	kill (GetServicePID (), kSIG_ReReadConfiguration);
#endif
}

void	Main::Pause ()
{
	AssertNotImplemented ();
}

void	Main::Continue ()
{
	AssertNotImplemented ();
}

Main::ServiceDescription	Main::GetServiceDescription () const
{
	return _sRep->GetServiceDescription ();
}

#if		qPlatform_POSIX
void	Main::SignalHandler (int signum)
{
	_sRep->SignalHandler (signum);
}
#endif

bool	Main::_HandleStandardCommandLineArgument (const String& arg)
{
	if (Execution::MatchesCommandLineArgument (arg, CommandNames::kRunAsService)) {
		RunAsService ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kStart)) {
		Start ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kStop)) {
		Stop ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kKill)) {
		Kill ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kRestart)) {
		Restart ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kReloadConfiguration)) {
		ReReadConfiguration ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kPause)) {
		Pause ();
		return true;
	}
	else if (Execution::MatchesCommandLineArgument (arg, CommandNames::kContinue)) {
		Continue ();
		return true;
	}
	/// MANY more neeeded, and fix to use named constants...
	return false;
}

<|endoftext|>


<commit_before>/* Copyright 2017 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.
==============================================================================*/
#include <utility>

#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/experimental/sql/driver_manager.h"
#include "tensorflow/core/kernels/data/experimental/sql/query_connection.h"
#include "tensorflow/core/lib/io/inputbuffer.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/strings/stringprintf.h"

namespace tensorflow {
namespace data {
namespace experimental {
namespace {

class SqlDatasetOp : public DatasetOpKernel {
 public:
  explicit SqlDatasetOp(OpKernelConstruction* ctx) : DatasetOpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
    for (const DataType& dt : output_types_) {
      OP_REQUIRES(ctx,
                  dt == DT_STRING || dt == DT_INT8 || dt == DT_INT16 ||
                      dt == DT_INT32 || dt == DT_INT64 || dt == DT_UINT8 ||
                      dt == DT_UINT16 || dt == DT_BOOL || dt == DT_DOUBLE,
                  errors::InvalidArgument(
                      "Each element of `output_types_` must be one of: "
                      "DT_STRING, DT_INT8, DT_INT16, DT_INT32, DT_INT64, "
                      "DT_UINT8, DT_UINT16, DT_BOOL, DT_DOUBLE "));
    }
    for (const PartialTensorShape& pts : output_shapes_) {
      OP_REQUIRES(ctx, pts.dims() == 0,
                  errors::InvalidArgument(
                      "Each element of `output_shapes_` must be a scalar."));
    }
  }
  void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override {
    tstring driver_name;
    OP_REQUIRES_OK(
        ctx, ParseScalarArgument<tstring>(ctx, "driver_name", &driver_name));

    tstring data_source_name;
    OP_REQUIRES_OK(ctx, ParseScalarArgument<tstring>(ctx, "data_source_name",
                                                     &data_source_name));

    tstring query;
    OP_REQUIRES_OK(ctx, ParseScalarArgument<tstring>(ctx, "query", &query));

    // TODO(b/64276826) Change this check when we add support for other
    // databases.
    OP_REQUIRES(ctx, driver_name == "sqlite",
                errors::InvalidArgument(tensorflow::strings::Printf(
                    "The database type, %s, is not supported by SqlDataset. "
                    "The set of supported databases is: {'sqlite'}.",
                    driver_name.c_str())));

    *output = new Dataset(ctx, driver_name, data_source_name, query,
                          output_types_, output_shapes_);
  }

 private:
  class Dataset : public DatasetBase {
   public:
    Dataset(OpKernelContext* ctx, const string& driver_name,
            const string& data_source_name, const string& query,
            const DataTypeVector& output_types,
            const std::vector<PartialTensorShape>& output_shapes)
        : DatasetBase(DatasetContext(ctx)),
          driver_name_(driver_name),
          data_source_name_(data_source_name),
          query_(query),
          output_types_(output_types),
          output_shapes_(output_shapes) {}

    std::unique_ptr<IteratorBase> MakeIteratorInternal(
        const string& prefix) const override {
      return absl::make_unique<Iterator>(
          Iterator::Params{this, strings::StrCat(prefix, "::Sql")});
    }

    const DataTypeVector& output_dtypes() const override {
      return output_types_;
    }

    const std::vector<PartialTensorShape>& output_shapes() const override {
      return output_shapes_;
    }

    string DebugString() const override { return "SqlDatasetOp::Dataset"; }

    Status CheckExternalState() const override { return Status::OK(); }

   protected:
    Status AsGraphDefInternal(SerializationContext* ctx,
                              DatasetGraphDefBuilder* b,
                              Node** output) const override {
      Node* driver_name_node;
      TF_RETURN_IF_ERROR(b->AddScalar(driver_name_, &driver_name_node));
      Node* data_source_name_node;
      TF_RETURN_IF_ERROR(
          b->AddScalar(data_source_name_, &data_source_name_node));
      Node* query_node;
      TF_RETURN_IF_ERROR(b->AddScalar(query_, &query_node));
      TF_RETURN_IF_ERROR(b->AddDataset(
          this, {driver_name_node, data_source_name_node, query_node}, output));
      return Status::OK();
    }

   private:
    class Iterator : public DatasetIterator<Dataset> {
     public:
      explicit Iterator(const Params& params)
          : DatasetIterator<Dataset>(params) {}
      ~Iterator() override {
        if (query_connection_initialized_) {
          Status s = query_connection_->Close();
          if (!s.ok()) {
            LOG(WARNING) << "Failed to close query connection: " << s;
          }
        }
      }

      Status GetNextInternal(IteratorContext* ctx,
                             std::vector<Tensor>* out_tensors,
                             bool* end_of_sequence) override {
        mutex_lock l(mu_);
        if (!query_connection_initialized_) {
          TF_RETURN_IF_ERROR(InitializeQueryConnection());
        }
        next_calls_++;
        return query_connection_->GetNext(ctx, out_tensors, end_of_sequence);
      }

     protected:
      std::shared_ptr<model::Node> CreateNode(
          IteratorContext* ctx, model::Node::Args args) const override {
        return model::MakeSourceNode(std::move(args));
      }

      Status SaveInternal(IteratorStateWriter* writer) override {
        mutex_lock l(mu_);
        if (query_connection_initialized_) {
          TF_RETURN_IF_ERROR(
              writer->WriteScalar(full_name("next_calls"), next_calls_));
        }
        return Status::OK();
      }

      Status RestoreInternal(IteratorContext* ctx,
                             IteratorStateReader* reader) override {
        mutex_lock l(mu_);
        if (reader->Contains(full_name("next_calls"))) {
          TF_RETURN_IF_ERROR(InitializeQueryConnection());
          TF_RETURN_IF_ERROR(
              reader->ReadScalar(full_name("next_calls"), &next_calls_));
          int64 rem_next_calls = next_calls_;
          std::vector<Tensor> out_tensors;
          bool end_of_sequence = false;
          while (rem_next_calls--) {
            TF_RETURN_IF_ERROR(query_connection_->GetNext(ctx, &out_tensors,
                                                          &end_of_sequence));
            out_tensors.clear();
          }
        } else {
          query_connection_initialized_ = false;
        }
        return Status::OK();
      }

     private:
      Status InitializeQueryConnection() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
        query_connection_initialized_ = true;
        query_connection_ =
            sql::DriverManager::CreateQueryConnection(dataset()->driver_name_);
        Status s = query_connection_->Open(dataset()->data_source_name_,
                                           dataset()->query_,
                                           dataset()->output_types_);
        next_calls_ = 0;
        if (!s.ok()) {
          LOG(WARNING) << "Failed to connect to database: " << s;
          return s;
        }
        return Status::OK();
      }

      mutex mu_;
      // TODO(b/129062371): explore ways to seek into a SQLite databases.
      int64 next_calls_ GUARDED_BY(mu_) = 0;
      std::unique_ptr<sql::QueryConnection> query_connection_ GUARDED_BY(mu_);
      bool query_connection_initialized_ GUARDED_BY(mu_) = false;
    };
    const tstring driver_name_;
    const tstring data_source_name_;
    const tstring query_;
    const DataTypeVector output_types_;
    const std::vector<PartialTensorShape> output_shapes_;
  };
  DataTypeVector output_types_;
  std::vector<PartialTensorShape> output_shapes_;
};

REGISTER_KERNEL_BUILDER(Name("SqlDataset").Device(DEVICE_CPU), SqlDatasetOp);
REGISTER_KERNEL_BUILDER(Name("ExperimentalSqlDataset").Device(DEVICE_CPU),
                        SqlDatasetOp);

}  // namespace
}  // namespace experimental
}  // namespace data
}  // namespace tensorflow
<commit_msg>Fix SqlDataset fails to raise StopIteration issue when combined with batch<commit_after>/* Copyright 2017 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.
==============================================================================*/
#include <utility>

#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/experimental/sql/driver_manager.h"
#include "tensorflow/core/kernels/data/experimental/sql/query_connection.h"
#include "tensorflow/core/lib/io/inputbuffer.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/strings/stringprintf.h"

namespace tensorflow {
namespace data {
namespace experimental {
namespace {

class SqlDatasetOp : public DatasetOpKernel {
 public:
  explicit SqlDatasetOp(OpKernelConstruction* ctx) : DatasetOpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
    for (const DataType& dt : output_types_) {
      OP_REQUIRES(ctx,
                  dt == DT_STRING || dt == DT_INT8 || dt == DT_INT16 ||
                      dt == DT_INT32 || dt == DT_INT64 || dt == DT_UINT8 ||
                      dt == DT_UINT16 || dt == DT_BOOL || dt == DT_DOUBLE,
                  errors::InvalidArgument(
                      "Each element of `output_types_` must be one of: "
                      "DT_STRING, DT_INT8, DT_INT16, DT_INT32, DT_INT64, "
                      "DT_UINT8, DT_UINT16, DT_BOOL, DT_DOUBLE "));
    }
    for (const PartialTensorShape& pts : output_shapes_) {
      OP_REQUIRES(ctx, pts.dims() == 0,
                  errors::InvalidArgument(
                      "Each element of `output_shapes_` must be a scalar."));
    }
  }
  void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override {
    tstring driver_name;
    OP_REQUIRES_OK(
        ctx, ParseScalarArgument<tstring>(ctx, "driver_name", &driver_name));

    tstring data_source_name;
    OP_REQUIRES_OK(ctx, ParseScalarArgument<tstring>(ctx, "data_source_name",
                                                     &data_source_name));

    tstring query;
    OP_REQUIRES_OK(ctx, ParseScalarArgument<tstring>(ctx, "query", &query));

    // TODO(b/64276826) Change this check when we add support for other
    // databases.
    OP_REQUIRES(ctx, driver_name == "sqlite",
                errors::InvalidArgument(tensorflow::strings::Printf(
                    "The database type, %s, is not supported by SqlDataset. "
                    "The set of supported databases is: {'sqlite'}.",
                    driver_name.c_str())));

    *output = new Dataset(ctx, driver_name, data_source_name, query,
                          output_types_, output_shapes_);
  }

 private:
  class Dataset : public DatasetBase {
   public:
    Dataset(OpKernelContext* ctx, const string& driver_name,
            const string& data_source_name, const string& query,
            const DataTypeVector& output_types,
            const std::vector<PartialTensorShape>& output_shapes)
        : DatasetBase(DatasetContext(ctx)),
          driver_name_(driver_name),
          data_source_name_(data_source_name),
          query_(query),
          output_types_(output_types),
          output_shapes_(output_shapes) {}

    std::unique_ptr<IteratorBase> MakeIteratorInternal(
        const string& prefix) const override {
      return absl::make_unique<Iterator>(
          Iterator::Params{this, strings::StrCat(prefix, "::Sql")});
    }

    const DataTypeVector& output_dtypes() const override {
      return output_types_;
    }

    const std::vector<PartialTensorShape>& output_shapes() const override {
      return output_shapes_;
    }

    string DebugString() const override { return "SqlDatasetOp::Dataset"; }

    Status CheckExternalState() const override { return Status::OK(); }

   protected:
    Status AsGraphDefInternal(SerializationContext* ctx,
                              DatasetGraphDefBuilder* b,
                              Node** output) const override {
      Node* driver_name_node;
      TF_RETURN_IF_ERROR(b->AddScalar(driver_name_, &driver_name_node));
      Node* data_source_name_node;
      TF_RETURN_IF_ERROR(
          b->AddScalar(data_source_name_, &data_source_name_node));
      Node* query_node;
      TF_RETURN_IF_ERROR(b->AddScalar(query_, &query_node));
      TF_RETURN_IF_ERROR(b->AddDataset(
          this, {driver_name_node, data_source_name_node, query_node}, output));
      return Status::OK();
    }

   private:
    class Iterator : public DatasetIterator<Dataset> {
     public:
      explicit Iterator(const Params& params)
          : DatasetIterator<Dataset>(params) {}
      ~Iterator() override {
        if (query_connection_initialized_) {
          Status s = query_connection_->Close();
          if (!s.ok()) {
            LOG(WARNING) << "Failed to close query connection: " << s;
          }
        }
      }

      Status GetNextInternal(IteratorContext* ctx,
                             std::vector<Tensor>* out_tensors,
                             bool* end_of_sequence) override {
        mutex_lock l(mu_);
        if (!query_connection_initialized_) {
          TF_RETURN_IF_ERROR(InitializeQueryConnection());
        }
	Status status = Status::OK();
	if (!end_of_sequence_) {
          next_calls_++;
          status = query_connection_->GetNext(ctx, out_tensors, &end_of_sequence_);
	}
	*end_of_sequence = end_of_sequence_;
	return status;
      }

     protected:
      std::shared_ptr<model::Node> CreateNode(
          IteratorContext* ctx, model::Node::Args args) const override {
        return model::MakeSourceNode(std::move(args));
      }

      Status SaveInternal(IteratorStateWriter* writer) override {
        mutex_lock l(mu_);
        if (query_connection_initialized_) {
          TF_RETURN_IF_ERROR(
              writer->WriteScalar(full_name("next_calls"), next_calls_));
        }
        return Status::OK();
      }

      Status RestoreInternal(IteratorContext* ctx,
                             IteratorStateReader* reader) override {
        mutex_lock l(mu_);
        if (reader->Contains(full_name("next_calls"))) {
          TF_RETURN_IF_ERROR(InitializeQueryConnection());
          TF_RETURN_IF_ERROR(
              reader->ReadScalar(full_name("next_calls"), &next_calls_));
          int64 rem_next_calls = next_calls_;
          std::vector<Tensor> out_tensors;
          end_of_sequence_ = false;
          while (rem_next_calls--) {
            TF_RETURN_IF_ERROR(query_connection_->GetNext(ctx, &out_tensors,
                                                          &end_of_sequence_));
            out_tensors.clear();
          }
        } else {
          query_connection_initialized_ = false;
          end_of_sequence_ = false;
        }
        return Status::OK();
      }

     private:
      Status InitializeQueryConnection() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
        query_connection_initialized_ = true;
        end_of_sequence_ = false;
        query_connection_ =
            sql::DriverManager::CreateQueryConnection(dataset()->driver_name_);
        Status s = query_connection_->Open(dataset()->data_source_name_,
                                           dataset()->query_,
                                           dataset()->output_types_);
        next_calls_ = 0;
        if (!s.ok()) {
          LOG(WARNING) << "Failed to connect to database: " << s;
          return s;
        }
        return Status::OK();
      }

      mutex mu_;
      // TODO(b/129062371): explore ways to seek into a SQLite databases.
      int64 next_calls_ GUARDED_BY(mu_) = 0;
      std::unique_ptr<sql::QueryConnection> query_connection_ GUARDED_BY(mu_);
      bool query_connection_initialized_ GUARDED_BY(mu_) = false;
      bool end_of_sequence_ GUARDED_BY(mu_) = false;
    };
    const tstring driver_name_;
    const tstring data_source_name_;
    const tstring query_;
    const DataTypeVector output_types_;
    const std::vector<PartialTensorShape> output_shapes_;
  };
  DataTypeVector output_types_;
  std::vector<PartialTensorShape> output_shapes_;
};

REGISTER_KERNEL_BUILDER(Name("SqlDataset").Device(DEVICE_CPU), SqlDatasetOp);
REGISTER_KERNEL_BUILDER(Name("ExperimentalSqlDataset").Device(DEVICE_CPU),
                        SqlDatasetOp);

}  // namespace
}  // namespace experimental
}  // namespace data
}  // namespace tensorflow
<|endoftext|>


<commit_before>#include "palabos3D.h"
#include "palabos3D.hh"
#include <vector>
#include <cmath>

using namespace plb;
using namespace plb::descriptors;
using namespace std;

typedef double T;
typedef Array<T,3> Velocity;
//#define DESCRIPTOR descriptors::D3Q27Descriptor
#define DESCRIPTOR MRTD3Q19Descriptor
 typedef MRTdynamics<T,DESCRIPTOR> BackgroundDynamics;
 typedef AnechoicMRTdynamics<T,DESCRIPTOR> AnechoicBackgroundDynamics;

// ---------------------------------------------
// Includes of acoustics resources
#include "acoustics/acoustics3D.h"
using namespace plb_acoustics_3D;
// ---------------------------------------------

const T rho0 = 1;
const T drho = rho0/10;

void writeGifs(MultiBlockLattice3D<T,DESCRIPTOR>& lattice, plint iter){
    const plint nx = lattice.getNx();
    const plint ny = lattice.getNy();
    const plint nz = lattice.getNz();

    const plint imSize = 600;
    ImageWriter<T> imageWriter("leeloo");

    Box3D slice(0, nx-1, 0, ny-1, nz/2, nz/2);
    //imageWriter.writeGif(createFileName("u", iT, 6),
     //*computeDensity(lattice), );

    imageWriter.writeGif( createFileName("rho", iter, 6),
    *computeDensity(lattice, slice),
    (T) rho0 - drho/1000000, (T) rho0 + drho/1000000, imSize, imSize);
}

void writeVTK(MultiBlockLattice3D<T,DESCRIPTOR>& lattice, plint iter){
        VtkImageOutput3D<T> vtkOut(createFileName("vtk", iter, 6), 1.);
        vtkOut.writeData<float>(*computeDensity(lattice), "density", 1.);
        vtkOut.writeData<3,float>(*computeVelocity(lattice), "velocity", 1.);
}

int main(int argc, char **argv){
    plbInit(&argc, &argv);
    std::string fNameOut = "tmp";

    const plint nx = 100;
    const plint ny = 100;
    const plint nz = 100;
    const T lattice_speed_sound = 1/sqrt(3);

    const T omega = 1.9;
    const plint maxT = 4000;

    const T lx =  2.3;
    const T dx =lx/nx;

    Array<T,3> u0(0, 0, 0);

    plint cx = util::roundToInt(nx/2);
    plint cy = util::roundToInt(ny/2);
    plint cz = util::roundToInt(nz/2);
    Array<T,3> centerLB(cx , cy, cz);

    global::directories().setOutputDir(fNameOut+"/");

    T rhoBar_target = 0;
    Array<T,3> j_target(0, 0, 0);
    T delta = 30;
    AnechoicBackgroundDynamics *anechoicDynamics =
                new AnechoicBackgroundDynamics(omega);
    anechoicDynamics->setDelta((T) delta);
    anechoicDynamics->setRhoBar_target(rhoBar_target);
    //j_target[0] = -j_target[0];
    anechoicDynamics->setJ_target(j_target);
    MultiBlockLattice3D<T, DESCRIPTOR> lattice(nx,ny,nz, anechoicDynamics);

    pcout << "Creation of the lattice." << endl;

    // Switch off periodicity.
    lattice.periodicity().toggleAll(false);

    pcout << "Initilization of rho and u." << endl;
    initializeAtEquilibrium( lattice, lattice.getBoundingBox(), rho0 , u0 );

    plint size_square = 80;
    Box3D square(
    nx/2 - size_square/2, nx/2 + size_square/2,
    ny/2 - size_square/2, ny/2 + size_square/2,
    nz/2 - size_square/2, nz/2 + size_square/2);
    defineDynamics(lattice, square, new BackgroundDynamics(omega));

    // Anechoic Condition
    /*T rhoBar_target = 0;
    Array<T,3> j_target(0, 0, 0);
    T size_anechoic_buffer = 30;
    // Define Anechoic Boards
    defineAnechoicBoards(nx, ny, nz, lattice, size_anechoic_buffer,
      omega, j_target, j_target, j_target, j_target, j_target, j_target,
      rhoBar_target);*/

    lattice.initialize();

    pcout << std::endl << "Voxelizing the domain." << std::endl;
    pcout << std::endl << "dx:" << dx << std::endl;

    pcout << "Simulation begins" << endl;

    plb_ofstream history_pressures("history_pressures.dat");
    for (plint iT=0; iT<maxT; ++iT){
        if (iT != 0){
            T lattice_speed_sound = 1/sqrt(3);
            T rho_changing = 1. + drho*sin(2*M_PI*(lattice_speed_sound/20)*iT);
            Box3D impulse(nx/2 + 20, nx/2 + 20, ny/2 + 20, ny/2 + 20, nz/2 + 20, nz/2 + 20);
            initializeAtEquilibrium( lattice, impulse, rho_changing, u0 );
        }

        if (iT % 100 == 0 && iT>0) {
            pcout << "Iteration " << iT << endl;
            //writeGifs(lattice,iT);
            //writeVTK(lattice, iT);
        }

        history_pressures << setprecision(10) << lattice.get(nx/2+30, ny/2+30, nz/2+30).computeDensity() - rho0 << endl;
        lattice.collideAndStream();

    }

    pcout << "End of simulation at iteration " << endl;

}
<commit_msg>Its more interesting to avoid bug of dynamics<commit_after>#include "palabos3D.h"
#include "palabos3D.hh"
#include <vector>
#include <cmath>

using namespace plb;
using namespace plb::descriptors;
using namespace std;

typedef double T;
typedef Array<T,3> Velocity;
//#define DESCRIPTOR descriptors::D3Q27Descriptor
#define DESCRIPTOR MRTD3Q19Descriptor
 typedef MRTdynamics<T,DESCRIPTOR> BackgroundDynamics;
 typedef AnechoicMRTdynamics<T,DESCRIPTOR> AnechoicBackgroundDynamics;

// ---------------------------------------------
// Includes of acoustics resources
#include "acoustics/acoustics3D.h"
using namespace plb_acoustics_3D;
// ---------------------------------------------

const T rho0 = 1;
const T drho = rho0/10;

void writeGifs(MultiBlockLattice3D<T,DESCRIPTOR>& lattice, plint iter){
    const plint nx = lattice.getNx();
    const plint ny = lattice.getNy();
    const plint nz = lattice.getNz();

    const plint imSize = 600;
    ImageWriter<T> imageWriter("leeloo");

    Box3D slice(0, nx-1, 0, ny-1, nz/2, nz/2);
    //imageWriter.writeGif(createFileName("u", iT, 6),
     //*computeDensity(lattice), );

    imageWriter.writeGif( createFileName("rho", iter, 6),
    *computeDensity(lattice, slice),
    (T) rho0 - drho/1000000, (T) rho0 + drho/1000000, imSize, imSize);
}

void writeVTK(MultiBlockLattice3D<T,DESCRIPTOR>& lattice, plint iter){
        VtkImageOutput3D<T> vtkOut(createFileName("vtk", iter, 6), 1.);
        vtkOut.writeData<float>(*computeDensity(lattice), "density", 1.);
        vtkOut.writeData<3,float>(*computeVelocity(lattice), "velocity", 1.);
}

int main(int argc, char **argv){
    plbInit(&argc, &argv);
    std::string fNameOut = "tmp";

    const plint nx = 100;
    const plint ny = 100;
    const plint nz = 100;
    const T lattice_speed_sound = 1/sqrt(3);

    const T omega = 1.9;
    const plint maxT = 4000;

    const T lx =  2.3;
    const T dx =lx/nx;

    Array<T,3> u0(0, 0, 0);

    plint cx = util::roundToInt(nx/2);
    plint cy = util::roundToInt(ny/2);
    plint cz = util::roundToInt(nz/2);
    Array<T,3> centerLB(cx , cy, cz);

    global::directories().setOutputDir(fNameOut+"/");

    T rhoBar_target = 0;
    Array<T,3> j_target(0, 0, 0);
    T delta = 30;
    AnechoicBackgroundDynamics *anechoicDynamics =
                new AnechoicBackgroundDynamics(omega);
    anechoicDynamics->setDelta((T) delta);
    anechoicDynamics->setRhoBar_target(rhoBar_target);
    //j_target[0] = -j_target[0];
    anechoicDynamics->setJ_target(j_target);
    MultiBlockLattice3D<T, DESCRIPTOR> lattice(nx, ny, nz, anechoicDynamics);
    defineDynamics(lattice, lattice.getBoundingBox(), new BackgroundDynamics(omega));

    pcout << "Creation of the lattice." << endl;

    // Switch off periodicity.
    lattice.periodicity().toggleAll(false);

    pcout << "Initilization of rho and u." << endl;
    initializeAtEquilibrium( lattice, lattice.getBoundingBox(), rho0 , u0 );

    plint size_square = 40;
    Box3D square(
    nx/2 - size_square/2, nx/2 + size_square/2,
    ny/2 - size_square/2, ny/2 + size_square/2,
    nz/2 - size_square/2, nz/2 + size_square/2);
    defineDynamics(lattice, square, anechoicDynamics);

    // Anechoic Condition
    /*T rhoBar_target = 0;
    Array<T,3> j_target(0, 0, 0);
    T size_anechoic_buffer = 30;
    // Define Anechoic Boards
    defineAnechoicBoards(nx, ny, nz, lattice, size_anechoic_buffer,
      omega, j_target, j_target, j_target, j_target, j_target, j_target,
      rhoBar_target);*/

    lattice.initialize();

    pcout << std::endl << "Voxelizing the domain." << std::endl;
    pcout << std::endl << "dx:" << dx << std::endl;

    pcout << "Simulation begins" << endl;

    plb_ofstream history_pressures("history_pressures.dat");
    for (plint iT=0; iT<maxT; ++iT){
        if (iT != 0){
            T lattice_speed_sound = 1/sqrt(3);
            T rho_changing = 1. + drho*sin(2*M_PI*(lattice_speed_sound/20)*iT);
            Box3D impulse(nx/2 + 20, nx/2 + 20, ny/2 + 20, ny/2 + 20, nz/2 + 20, nz/2 + 20);
            initializeAtEquilibrium( lattice, impulse, rho_changing, u0 );
        }

        if (iT % 100 == 0 && iT>0) {
            pcout << "Iteration " << iT << endl;
            //writeGifs(lattice,iT);
            //writeVTK(lattice, iT);
        }

        history_pressures << setprecision(10) << lattice.get(nx/2+30, ny/2+30, nz/2+30).computeDensity() - rho0 << endl;
        lattice.collideAndStream();

    }

    pcout << "End of simulation at iteration " << endl;

}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 <chrono>
#include <condition_variable>
#include <deque>
#include <memory>
#include <mutex>
#include <random>
#include <thread>

#include <folly/Synchronized.h>
#include <folly/portability/GMock.h>
#include <folly/portability/GTest.h>
#include <folly/synchronization/Baton.h>
#include <thrift/lib/cpp/concurrency/FunctionRunner.h>
#include <thrift/lib/cpp/concurrency/SFQThreadManager.h>

using namespace apache::thrift::concurrency;
using testing::_;
using testing::AnyNumber;
using testing::AtLeast;

class SFQThreadManagerTest : public testing::Test {
 public:
  MOCK_METHOD1(bogusTask, void(int));

 protected:
  std::shared_ptr<ThreadManager> newSFQTM(
      std::chrono::seconds perturb, size_t numQueues) {
    SFQThreadManagerConfig config;
    config.setPerturbInterval(perturb)
        .setNumFairQueuesForUpstream(numQueues)
        .setExecutors(
            {ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1)});
    return std::make_shared<SFQThreadManager>(std::move(config));
  }
};

// Verify tasks are executed at all.
TEST_F(SFQThreadManagerTest, SmokeTest) {
  auto tm = newSFQTM(std::chrono::seconds(1), 1);
  tm->start();
  ThreadManager::ExecutionScope es(PRIORITY::NORMAL);
  es.setTenantId(123);
  auto ka = tm->getKeepAlive(std::move(es), ThreadManager::Source::UPSTREAM);

  EXPECT_CALL(*this, bogusTask(0)).Times(1);
  ka->add([this]() { this->bogusTask(0); });
}

// Ensure the queuing is fair and that higher priority tasks pre-empt low pri.
TEST_F(SFQThreadManagerTest, FairnessPreemptTest) {
  // Disabling perturbation so we can actually test this.
  auto tm = newSFQTM(std::chrono::seconds(0), 10000);
  const auto source = ThreadManager::Source::UPSTREAM;
  tm->start();

  // This will dictate the expected order of placing the tasks.
  std::vector<folly::Baton<>> addOrderBaton(4);

  std::vector<folly::Baton<>> c0Baton(2), c1Baton(2);
  size_t c0{0}, c1{0};

  ThreadManager::ExecutionScope es(PRIORITY::NORMAL);
  es.setTenantId(0);
  tm->getKeepAlive(es, source)->add([&]() {
    addOrderBaton[0].wait();
    ++c0;
    c0Baton[0].post();
  });

  es.setTenantId(0);
  tm->getKeepAlive(es, source)->add([&]() {
    addOrderBaton[1].wait();
    ++c0;
    c0Baton[1].post();
  });

  es.setTenantId(1);
  tm->getKeepAlive(es, source)->add([&]() {
    addOrderBaton[2].wait();
    ++c1;
    c1Baton[0].post();
  });

  es.setTenantId(1);
  tm->getKeepAlive(es, source)->add([&]() {
    addOrderBaton[3].wait();
    ++c1;
    c1Baton[1].post();
  });

  // No tasks have run at this point.
  EXPECT_EQ(0, c0);
  EXPECT_EQ(0, c1);

  // Tenant 0 was added first, so we expect this to execute.
  addOrderBaton[0].post();
  c0Baton[0].wait();
  EXPECT_EQ(1, c0);
  EXPECT_EQ(0, c1);

  // Tenant 1 should be next even though it was added 3rd. Posting the 3rd
  // add-order baton would lock up here if it were unfair.
  addOrderBaton[2].post();
  c1Baton[0].wait();
  EXPECT_EQ(1, c0);
  EXPECT_EQ(1, c1);

  // Tenant 0 will then be up next. It was the task added 2nd.
  addOrderBaton[1].post();
  c0Baton[1].wait();
  EXPECT_EQ(2, c0);
  EXPECT_EQ(1, c1);

  // Tenant 1 would be up next, but let's preempt all this with a higher
  // priority source.
  folly::Baton<> hpribaton, hpribatonOuter;
  es = ThreadManager::ExecutionScope(PRIORITY::HIGH);
  es.setTenantId(123);
  tm->getKeepAlive(es, ThreadManager::Source::INTERNAL)->add([&]() {
    hpribaton.wait();
    hpribatonOuter.post();
  });
  hpribaton.post();
  hpribatonOuter.wait();

  // Now we should be able to execute the tenant 1 task after the
  // source-preempted task.
  addOrderBaton[3].post();
  c1Baton[1].wait();
  EXPECT_EQ(2, c0);
  EXPECT_EQ(2, c1);

  addOrderBaton.clear();
  c0Baton.clear();
  c1Baton.clear();
}
<commit_msg>Instrument SFQThreadManagerTest<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 <chrono>
#include <condition_variable>
#include <deque>
#include <memory>
#include <mutex>
#include <random>
#include <thread>

#include <folly/Synchronized.h>
#include <folly/portability/GMock.h>
#include <folly/portability/GTest.h>
#include <folly/synchronization/Baton.h>
#include <thrift/lib/cpp/concurrency/FunctionRunner.h>
#include <thrift/lib/cpp/concurrency/SFQThreadManager.h>

using namespace apache::thrift::concurrency;
using namespace std::literals::chrono_literals;
using testing::_;
using testing::AnyNumber;
using testing::AtLeast;

class SFQThreadManagerTest : public testing::Test {
 public:
  MOCK_METHOD1(bogusTask, void(int));

 protected:
  std::shared_ptr<ThreadManager> newSFQTM(
      std::chrono::seconds perturb, size_t numQueues) {
    SFQThreadManagerConfig config;
    config.setPerturbInterval(perturb)
        .setNumFairQueuesForUpstream(numQueues)
        .setExecutors(
            {ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1),
             ThreadManager::newSimpleThreadManager(1)});
    return std::make_shared<SFQThreadManager>(std::move(config));
  }
};

// Verify tasks are executed at all.
TEST_F(SFQThreadManagerTest, SmokeTest) {
  auto tm = newSFQTM(std::chrono::seconds(1), 1);
  tm->start();
  ThreadManager::ExecutionScope es(PRIORITY::NORMAL);
  es.setTenantId(123);
  auto ka = tm->getKeepAlive(std::move(es), ThreadManager::Source::UPSTREAM);

  EXPECT_CALL(*this, bogusTask(0)).Times(1);
  ka->add([this]() { this->bogusTask(0); });
}

// Ensure the queuing is fair and that higher priority tasks pre-empt low pri.
TEST_F(SFQThreadManagerTest, FairnessPreemptTest) {
  // Disabling perturbation so we can actually test this.
  auto tm = newSFQTM(std::chrono::seconds(0), 10000);
  const auto source = ThreadManager::Source::UPSTREAM;
  tm->start();

  // This will dictate the expected order of placing the tasks.
  std::vector<folly::Baton<>> addOrderBaton(4);

  std::vector<folly::Baton<>> c0Baton(2), c1Baton(2);
  size_t c0{0}, c1{0};

  ThreadManager::ExecutionScope es(PRIORITY::NORMAL);
  es.setTenantId(0);
  tm->getKeepAlive(es, source)->add([&]() {
    ASSERT_TRUE(addOrderBaton[0].try_wait_for(3s));
    ++c0;
    c0Baton[0].post();
  });

  es.setTenantId(0);
  tm->getKeepAlive(es, source)->add([&]() {
    ASSERT_TRUE(addOrderBaton[1].try_wait_for(3s));
    ++c0;
    c0Baton[1].post();
  });

  es.setTenantId(1);
  tm->getKeepAlive(es, source)->add([&]() {
    ASSERT_TRUE(addOrderBaton[2].try_wait_for(3s));
    ++c1;
    c1Baton[0].post();
  });

  es.setTenantId(1);
  tm->getKeepAlive(es, source)->add([&]() {
    ASSERT_TRUE(addOrderBaton[3].try_wait_for(3s));
    ++c1;
    c1Baton[1].post();
  });

  // No tasks have run at this point.
  EXPECT_EQ(0, c0);
  EXPECT_EQ(0, c1);

  // Tenant 0 was added first, so we expect this to execute.
  addOrderBaton[0].post();
  ASSERT_TRUE(c0Baton[0].try_wait_for(3s));
  EXPECT_EQ(1, c0);
  EXPECT_EQ(0, c1);

  // Tenant 1 should be next even though it was added 3rd. Posting the 3rd
  // add-order baton would lock up here if it were unfair.
  addOrderBaton[2].post();
  ASSERT_TRUE(c1Baton[0].try_wait_for(3s));
  EXPECT_EQ(1, c0);
  EXPECT_EQ(1, c1);

  // Tenant 0 will then be up next. It was the task added 2nd.
  addOrderBaton[1].post();
  ASSERT_TRUE(c0Baton[1].try_wait_for(3s));
  EXPECT_EQ(2, c0);
  EXPECT_EQ(1, c1);

  // Tenant 1 would be up next, but let's preempt all this with a higher
  // priority source.
  folly::Baton<> hpribaton, hpribatonOuter;
  es = ThreadManager::ExecutionScope(PRIORITY::HIGH);
  es.setTenantId(123);
  tm->getKeepAlive(es, ThreadManager::Source::INTERNAL)->add([&]() {
    ASSERT_TRUE(hpribaton.try_wait_for(3s));
    hpribatonOuter.post();
  });
  hpribaton.post();
  ASSERT_TRUE(hpribatonOuter.try_wait_for(3s));

  // Now we should be able to execute the tenant 1 task after the
  // source-preempted task.
  addOrderBaton[3].post();
  ASSERT_TRUE(c1Baton[1].try_wait_for(3s));
  EXPECT_EQ(2, c0);
  EXPECT_EQ(2, c1);

  addOrderBaton.clear();
  c0Baton.clear();
  c1Baton.clear();
}
<|endoftext|>