<commit_before>/* Copyright (c) Citrix Systems Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms,
 * with or without modification, are permitted provided
 * that the following conditions are met:
 *
 * *   Redistributions of source code must retain the above
 *     copyright notice, this list of conditions and the
 *     following disclaimer.
 * *   Redistributions 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.
 *
 * 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 COPYRIGHT HOLDER 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.
 */

#include "Setup.h"
#include "shellapi.h"
#include "shlobj.h"
#include "shlwapi.h"
#include "stdint.h"
#include <windows.h>
#include "setupapi.h"
#include "Sddl.h"

TCHAR * _vatallocprintf(const TCHAR *format, va_list args)
{
	int chars;
	chars = _vsctprintf(format, args);
	size_t size= sizeof(TCHAR) * (chars+1);
	TCHAR *space = (TCHAR *)malloc(size);
	if (space == NULL)
		return NULL;
	_vsntprintf_s(space, chars+1, chars, format, args);
	return space;
}

TCHAR * _tallocprintf(const TCHAR *format, ...)
{
	TCHAR *space;
	va_list args;
    va_start (args, format);
	space=_vatallocprintf(format, args);
	va_end(args);
	return space;
}

#include "brandcontrol.h"

void ErrMsg(const TCHAR *format, ...)
{
	va_list args;
    va_start (args, format);

	TCHAR* space = _vatallocprintf(format, args);
	if (space == NULL)
		goto fail1;

	va_end(args);
	MessageBox(NULL, space, getBrandingString(BRANDING_setupErr), MB_OK);
	free(space);
	return;
fail1:
	va_end(args);
}

bool runProcess(TCHAR* cmdLine, DWORD* exitcode)
{
	PROCESS_INFORMATION pi;
	STARTUPINFO si;
	memset(&si, 0, sizeof(STARTUPINFO));

	OutputDebugString(cmdLine);

	if (!CreateProcess(NULL, cmdLine, NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi)) {
		ErrMsg(getBrandingString(BRANDING_processFail),cmdLine, GetLastError());
		return false;
	}

	WaitForSingleObject(pi.hProcess, INFINITE);
	GetExitCodeProcess(pi.hProcess, exitcode);
	return true;
}

typedef struct {
	bool test;
	bool passive;
	bool quiet;
	bool norestart;
	bool forcerestart;
	bool legacy;
} arguments;

bool parseCommandLine(arguments* args)
{

	memset(args, 0, sizeof(arguments));

	int argCount;
	LPWSTR cli = GetCommandLineW();
	LPWSTR *szArgList = CommandLineToArgvW(cli, &argCount);
	for (int i=1; i<argCount ; i++) {
		if (!wcsncmp(szArgList[i],L"/TEST",sizeof(L"/TEST"))) {
			args->test = true;
		}
		else if (!wcsncmp(szArgList[i],L"/passive",sizeof(L"/passive"))) {
			args->passive = true;
		}
		else if (!wcsncmp(szArgList[i],L"/quiet",sizeof(L"/quiet"))) {
			args->quiet = true;
		}
		else if (!wcsncmp(szArgList[i],L"/norestart",sizeof(L"/norestart"))) {
			args->norestart = true;
		}
		else if (!wcsncmp(szArgList[i],L"/forcerestart",sizeof(L"/forcerestart"))) {
			args->forcerestart = true;
		}
		else {
			ErrMsg(getBrandingString(BRANDING_setupHelp));
			return false;
		}
	}

	return true;
}

const TCHAR* ma64 = getBrandingString(FILENAME_managementx64);
const TCHAR* ma32 = getBrandingString(FILENAME_managementx86);
const TCHAR* iw = _T("installwizard.msi");

const TCHAR* xenlegacy = getBrandingString(FILENAME_legacy);
const TCHAR* uninstallerfix = getBrandingString(FILENAME_legacyuninstallerfix);
const TCHAR* INSTALL_AGENT_REG_KEY = getBrandingString(BRANDING_installAgentRegKey);

TCHAR sysdir[MAX_PATH];
TCHAR workfile[MAX_PATH];
TCHAR logfile[MAX_PATH];
TCHAR* msiexec;

int getFileLocations()
{
	if (!GetSystemDirectory(sysdir, MAX_PATH)) {
		ErrMsg(getBrandingString(BRANDING_noSystemDir));
		return 0;
	}

	msiexec = _tallocprintf(_T("%s\\msiexec.exe"), sysdir);
	if ( msiexec == NULL) {
		ErrMsg(getBrandingString(BRANDING_errMSINoMem));
		return 0;
	}

	if (GetModuleFileName(NULL, workfile, MAX_PATH)>=MAX_PATH) {
		ErrMsg(getBrandingString(BRANDING_errFilePathNoMem));
		return 0;
	}
	PathRemoveFileSpec(workfile);

	if (SUCCEEDED(SHGetFolderPath(NULL, CSIDL_COMMON_APPDATA | CSIDL_FLAG_CREATE, NULL, 0 , logfile))) {
		PathAppend(logfile, (TCHAR *)getBrandingString(BRANDING_manufacturer));
		CreateDirectory(logfile, NULL);
		PathAppend(logfile, getBrandingString(BRANDING_setupLogDir));
		CreateDirectory(logfile, NULL);
		PathAppend(logfile,_T("Install.log"));
	}
	else {
		ErrMsg(getBrandingString(BRANDING_errNoLogPath));
		return 0;
	}
	return 1;
}

const TCHAR* getInstallMsiName(arguments* args)
{
	if (!args->test) {
		BOOL wow64;
		if (IsWow64Process(GetCurrentProcess(), &wow64)) {
			if (wow64) {
				return ma64;
			}
			else {
				return ma32;
			}
		}
		else {
			return ma32;
		}
	}
	else {
		args->forcerestart = false;  //Hack because old installer doesn't like forcereboot
		return iw;
	}
}

DWORD installMsi(arguments* args)
{
	DWORD exitcode;

	const TCHAR* installname = getInstallMsiName(args);
	TCHAR* cmdline = _tallocprintf(_T("\"%s\" /i\"%s\\%s\"%s%s /liwearucmopvx+! \"%s\""),
		msiexec,
		workfile,
		installname,
		( args->passive?_T(" /passive"): (args->quiet?_T(" /quiet"):_T(""))),
		( args->norestart?_T(" /norestart"):(args->forcerestart?_T(" /forcerestart"):_T(""))),
		logfile);
	if (cmdline == NULL) {
		ErrMsg(getBrandingString(BRANDING_errCmdLineNoMem));
		return 0;
	}

	runProcess(cmdline, &exitcode);

	return exitcode;
}

/*
 * WARNING: Call this function ONLY BEFORE installMsi()
 *          (and preferably save the value for later use)
 */
BOOL Installing(void)
{
	HKEY regKey;

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			INSTALL_AGENT_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&regKey) != ERROR_SUCCESS) {
		RegCloseKey(regKey);
		return TRUE; /* installing */
	} else {
		return FALSE; /* uninstalling */
	}

}

void waitForStatus() {
	HKEY  insAgntKey;
	HKEY  vmStateKey;
	TCHAR status[256];
	DWORD statusSize = 256;
	DWORD pvToolsOnFirstRun;
	DWORD pvToolsVerSize = sizeof(pvToolsOnFirstRun);

	TCHAR *VM_STATE_REG_KEY = _tallocprintf(
		_T("%s\\%s"),
		INSTALL_AGENT_REG_KEY,
		_T("VMState")
	);

	if (VM_STATE_REG_KEY == NULL) {
		goto fail1;
	}

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			INSTALL_AGENT_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&insAgntKey) != ERROR_SUCCESS) {
		goto fail2;
	}

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			VM_STATE_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&vmStateKey) != ERROR_SUCCESS) {
		goto fail3;
	}

	/*
	 * Gets populated in
	 * InstallAgent static ctor => State.VM static ctor
	 */
	if (RegQueryValueEx(
			vmStateKey,
			_T("PVToolsVersionOnFirstRun"),
			NULL,
			NULL,
			(LPBYTE)&pvToolsOnFirstRun,
			&pvToolsVerSize) != ERROR_SUCCESS) {
		goto fail4;
	}

	for(;;) {
		if (RegQueryValueEx(
				insAgntKey,
				_T("InstallStatus"),
				NULL,
				NULL,
				(LPBYTE)status,
				&statusSize) != ERROR_SUCCESS ||
		    _tcsicmp(status, _T("Installing")) == 0) {
			goto needcontinue;
		}

		if (_tcsicmp(status, _T("NeedsReboot")) == 0) {
			/*
			 * If PV drivers other than the latest
			 * (currently 8.x) are found, we should prompt
			 * for reboot after removing them from
			 * filters and disabling their bootstart
			 */
			if (pvToolsOnFirstRun == 7) {
				SetupPromptReboot(NULL, NULL, FALSE);
			}
			goto done;
		}

needcontinue:
		SleepEx(2000, TRUE);
		continue;
done:
fail4:
		RegCloseKey(vmStateKey);
fail3:
		RegCloseKey(insAgntKey);
fail2:
		free(VM_STATE_REG_KEY);
fail1:
		return;
	}
}

#define dotnet4 _T("Software\\Microsoft\\NET Framework Setup\\NDP\\v4\\Client")
#define dotnet35 _T("Software\\Microsoft\\NET Framework Setup\\NDP\\v3.5")
bool checkDotNet() {
	HKEY outkey;

	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, dotnet4, 0, KEY_READ, &outkey)!=ERROR_SUCCESS) {
		goto check35;
	}
	if (RegQueryValueEx(outkey, _T("Install"), NULL, NULL, NULL, NULL)==ERROR_SUCCESS) {
		RegCloseKey(outkey);
		return true;
	}
	RegCloseKey(outkey);
check35:
	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, dotnet35, 0, KEY_READ, &outkey)!=ERROR_SUCCESS) {
		return false;
	}
	if (RegQueryValueEx(outkey, _T("Install"), NULL, NULL, NULL, NULL)==ERROR_SUCCESS) {
		RegCloseKey(outkey);
		return true;
	}
	RegCloseKey(outkey);
	return false;
}

int APIENTRY _tWinMain(_In_ HINSTANCE hInstance,
                     _In_opt_ HINSTANCE hPrevInstance,
                     _In_ LPTSTR    lpCmdLine,
                     _In_ int       nCmdShow)
{
	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);
	DWORD msiResult;

	// Parse command line
	// Acceptable options:
	// /TEST - test managementagent(x86/x64).msi based install
	// /quiet - silent, noninteractive install
	// /passive - passive install
	// /norestart - does not restart vm
	// /forcerestart - will always restart vm

	arguments args;

	BOOL installing = Installing();

	if (!parseCommandLine(&args))
		return 0;

	// We require .net to run

	if (!checkDotNet()) {
		ErrMsg(getBrandingString(BRANDING_errDotNetNeeded));
		return 0;
	}

	//Locate Files

	if (!getFileLocations()) {
		return 0;
	}

	msiResult = installMsi(&args);

	if (msiResult != ERROR_SUCCESS &&
			msiResult != ERROR_SUCCESS_REBOOT_INITIATED &&
			msiResult != ERROR_SUCCESS_REBOOT_REQUIRED) {
		ErrMsg(
			getBrandingString(BRANDING_errMSIInstallFail),
			msiResult,
			logfile
		);
		return msiResult;
	}

	if (installing) {
		waitForStatus();
	}

	return 0;
}

<commit_msg>CP-17859 : Setup.exe passes autoupdate parameters<commit_after>/* Copyright (c) Citrix Systems Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms,
 * with or without modification, are permitted provided
 * that the following conditions are met:
 *
 * *   Redistributions of source code must retain the above
 *     copyright notice, this list of conditions and the
 *     following disclaimer.
 * *   Redistributions 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.
 *
 * 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 COPYRIGHT HOLDER 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.
 */

#include "Setup.h"
#include "shellapi.h"
#include "shlobj.h"
#include "shlwapi.h"
#include "stdint.h"
#include <windows.h>
#include "setupapi.h"
#include "Sddl.h"

TCHAR * _vatallocprintf(const TCHAR *format, va_list args)
{
	int chars;
	chars = _vsctprintf(format, args);
	size_t size= sizeof(TCHAR) * (chars+1);
	TCHAR *space = (TCHAR *)malloc(size);
	if (space == NULL)
		return NULL;
	_vsntprintf_s(space, chars+1, chars, format, args);
	return space;
}

TCHAR * _tallocprintf(const TCHAR *format, ...)
{
	TCHAR *space;
	va_list args;
    va_start (args, format);
	space=_vatallocprintf(format, args);
	va_end(args);
	return space;
}

#include "brandcontrol.h"

void ErrMsg(const TCHAR *format, ...)
{
	va_list args;
    va_start (args, format);

	TCHAR* space = _vatallocprintf(format, args);
	if (space == NULL)
		goto fail1;

	va_end(args);
	MessageBox(NULL, space, getBrandingString(BRANDING_setupErr), MB_OK);
	free(space);
	return;
fail1:
	va_end(args);
}

bool runProcess(TCHAR* cmdLine, DWORD* exitcode)
{
	PROCESS_INFORMATION pi;
	STARTUPINFO si;
	memset(&si, 0, sizeof(STARTUPINFO));

	OutputDebugString(cmdLine);

	if (!CreateProcess(NULL, cmdLine, NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi)) {
		ErrMsg(getBrandingString(BRANDING_processFail),cmdLine, GetLastError());
		return false;
	}

	WaitForSingleObject(pi.hProcess, INFINITE);
	GetExitCodeProcess(pi.hProcess, exitcode);
	return true;
}

typedef struct {
	bool test;
	bool passive;
	bool quiet;
	bool norestart;
	bool forcerestart;
	bool legacy;
    int  driverupdate;
    int  autoupdate;
    int  driverinstall;
} arguments;

bool parseCommandLine(arguments* args)
{

	memset(args, 0, sizeof(arguments));

	int argCount;
	LPWSTR cli = GetCommandLineW();
	LPWSTR *szArgList = CommandLineToArgvW(cli, &argCount);
    args->driverinstall=-1;
    args->autoupdate=-1;
    args->driverupdate=-1;
	for (int i=1; i<argCount ; i++) {
		if (!wcsncmp(szArgList[i],L"/TEST",sizeof(L"/TEST"))) {
			args->test = true;
		}
		else if (!wcsncmp(szArgList[i],L"/passive",sizeof(L"/passive"))) {
			args->passive = true;
		}
		else if (!wcsncmp(szArgList[i],L"/quiet",sizeof(L"/quiet"))) {
			args->quiet = true;
		}
		else if (!wcsncmp(szArgList[i],L"/norestart",sizeof(L"/norestart"))) {
			args->norestart = true;
		}
		else if (!wcsncmp(szArgList[i],L"/forcerestart",sizeof(L"/forcerestart"))) {
			args->forcerestart = true;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWDRIVERINSTALL=NO",sizeof(L"ALLOWDRIVERINSTALL=NO"))) {
			args->driverinstall=0;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWDRIVERINSTALL=YES",sizeof(L"ALLOWDRIVERINSTALL=YES"))) {
			args->driverinstall=1;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWAUTOUPDATE=NO",sizeof(L"ALLOWAUTOUPDATE=NO"))) {
			args->autoupdate=0;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWAUTOUPDATE=YES",sizeof(L"ALLOWAUTOUPDATE=YES"))) {
			args->autoupdate=1;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWDRIVERUPDATE=NO",sizeof(L"ALLOWDRIVERUPDATE=NO"))) {
			args->driverupdate=0;
		}
		else if (!wcsncmp(szArgList[i],L"ALLOWDRIVERUPDATE=YES",sizeof(L"ALLOWDRIVERUPDATE=YES"))) {
			args->driverupdate=1;
		}
		else {
			ErrMsg(getBrandingString(BRANDING_setupHelp));
			return false;
		}
	}

	return true;
}

const TCHAR* ma64 = getBrandingString(FILENAME_managementx64);
const TCHAR* ma32 = getBrandingString(FILENAME_managementx86);
const TCHAR* iw = _T("installwizard.msi");

const TCHAR* xenlegacy = getBrandingString(FILENAME_legacy);
const TCHAR* uninstallerfix = getBrandingString(FILENAME_legacyuninstallerfix);
const TCHAR* INSTALL_AGENT_REG_KEY = getBrandingString(BRANDING_installAgentRegKey);

TCHAR sysdir[MAX_PATH];
TCHAR workfile[MAX_PATH];
TCHAR logfile[MAX_PATH];
TCHAR* msiexec;

int getFileLocations()
{
	if (!GetSystemDirectory(sysdir, MAX_PATH)) {
		ErrMsg(getBrandingString(BRANDING_noSystemDir));
		return 0;
	}

	msiexec = _tallocprintf(_T("%s\\msiexec.exe"), sysdir);
	if ( msiexec == NULL) {
		ErrMsg(getBrandingString(BRANDING_errMSINoMem));
		return 0;
	}

	if (GetModuleFileName(NULL, workfile, MAX_PATH)>=MAX_PATH) {
		ErrMsg(getBrandingString(BRANDING_errFilePathNoMem));
		return 0;
	}
	PathRemoveFileSpec(workfile);

	if (SUCCEEDED(SHGetFolderPath(NULL, CSIDL_COMMON_APPDATA | CSIDL_FLAG_CREATE, NULL, 0 , logfile))) {
		PathAppend(logfile, (TCHAR *)getBrandingString(BRANDING_manufacturer));
		CreateDirectory(logfile, NULL);
		PathAppend(logfile, getBrandingString(BRANDING_setupLogDir));
		CreateDirectory(logfile, NULL);
		PathAppend(logfile,_T("Install.log"));
	}
	else {
		ErrMsg(getBrandingString(BRANDING_errNoLogPath));
		return 0;
	}
	return 1;
}

const TCHAR* getInstallMsiName(arguments* args)
{
	if (!args->test) {
		BOOL wow64;
		if (IsWow64Process(GetCurrentProcess(), &wow64)) {
			if (wow64) {
				return ma64;
			}
			else {
				return ma32;
			}
		}
		else {
			return ma32;
		}
	}
	else {
		args->forcerestart = false;  //Hack because old installer doesn't like forcereboot
		return iw;
	}
}

DWORD installMsi(arguments* args)
{
	DWORD exitcode;

	const TCHAR* installname = getInstallMsiName(args);
	TCHAR* cmdline = _tallocprintf(_T("\"%s\" /i\"%s\\%s\"%s%s%s%s%s /liwearucmopvx+! \"%s\""),
		msiexec,
		workfile,
		installname,
		( args->passive?_T(" /passive"): (args->quiet?_T(" /quiet"):_T(""))),
		( args->norestart?_T(" /norestart"):(args->forcerestart?_T(" /forcerestart"):_T(""))),
        ( (args->autoupdate >= 0)?
                (args->autoupdate?_T(" ALLOWAUTOUPDATE=YES"):_T(" ALLOWAUTOUPDATE=NO")):_T("")),
        ( (args->driverupdate >= 0)?
                (args->driverupdate?_T(" ALLOWDRIVERUPDATE=YES"):_T(" ALLOWDRIVERUPDATE=NO")):_T("")),
        ( (args->driverinstall >= 0)?
                (args->driverinstall?_T(" ALLOWDRIVERINSTALL=YES"):_T(" ALLOWDRIVERINSTALL=NO")):_T("")),
		logfile);
	if (cmdline == NULL) {
		ErrMsg(getBrandingString(BRANDING_errCmdLineNoMem));
		return 0;
	}

	runProcess(cmdline, &exitcode);

	return exitcode;
}

/*
 * WARNING: Call this function ONLY BEFORE installMsi()
 *          (and preferably save the value for later use)
 */
BOOL Installing(void)
{
	HKEY regKey;

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			INSTALL_AGENT_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&regKey) != ERROR_SUCCESS) {
		RegCloseKey(regKey);
		return TRUE; /* installing */
	} else {
		return FALSE; /* uninstalling */
	}

}

void waitForStatus() {
	HKEY  insAgntKey;
	HKEY  vmStateKey;
	TCHAR status[256];
	DWORD statusSize = 256;
	DWORD pvToolsOnFirstRun;
	DWORD pvToolsVerSize = sizeof(pvToolsOnFirstRun);

	TCHAR *VM_STATE_REG_KEY = _tallocprintf(
		_T("%s\\%s"),
		INSTALL_AGENT_REG_KEY,
		_T("VMState")
	);

	if (VM_STATE_REG_KEY == NULL) {
		goto fail1;
	}

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			INSTALL_AGENT_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&insAgntKey) != ERROR_SUCCESS) {
		goto fail2;
	}

	if (RegOpenKeyEx(
			HKEY_LOCAL_MACHINE,
			VM_STATE_REG_KEY,
			0,
			KEY_READ | KEY_WOW64_64KEY,
			&vmStateKey) != ERROR_SUCCESS) {
		goto fail3;
	}

	/*
	 * Gets populated in
	 * InstallAgent static ctor => State.VM static ctor
	 */
	if (RegQueryValueEx(
			vmStateKey,
			_T("PVToolsVersionOnFirstRun"),
			NULL,
			NULL,
			(LPBYTE)&pvToolsOnFirstRun,
			&pvToolsVerSize) != ERROR_SUCCESS) {
		goto fail4;
	}

	for(;;) {
		if (RegQueryValueEx(
				insAgntKey,
				_T("InstallStatus"),
				NULL,
				NULL,
				(LPBYTE)status,
				&statusSize) != ERROR_SUCCESS ||
		    _tcsicmp(status, _T("Installing")) == 0) {
			goto needcontinue;
		}

		if (_tcsicmp(status, _T("NeedsReboot")) == 0) {
			/*
			 * If PV drivers other than the latest
			 * (currently 8.x) are found, we should prompt
			 * for reboot after removing them from
			 * filters and disabling their bootstart
			 */
			if (pvToolsOnFirstRun == 7) {
				SetupPromptReboot(NULL, NULL, FALSE);
			}
			goto done;
		}

needcontinue:
		SleepEx(2000, TRUE);
		continue;
done:
fail4:
		RegCloseKey(vmStateKey);
fail3:
		RegCloseKey(insAgntKey);
fail2:
		free(VM_STATE_REG_KEY);
fail1:
		return;
	}
}

#define dotnet4 _T("Software\\Microsoft\\NET Framework Setup\\NDP\\v4\\Client")
#define dotnet35 _T("Software\\Microsoft\\NET Framework Setup\\NDP\\v3.5")
bool checkDotNet() {
	HKEY outkey;

	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, dotnet4, 0, KEY_READ, &outkey)!=ERROR_SUCCESS) {
		goto check35;
	}
	if (RegQueryValueEx(outkey, _T("Install"), NULL, NULL, NULL, NULL)==ERROR_SUCCESS) {
		RegCloseKey(outkey);
		return true;
	}
	RegCloseKey(outkey);
check35:
	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, dotnet35, 0, KEY_READ, &outkey)!=ERROR_SUCCESS) {
		return false;
	}
	if (RegQueryValueEx(outkey, _T("Install"), NULL, NULL, NULL, NULL)==ERROR_SUCCESS) {
		RegCloseKey(outkey);
		return true;
	}
	RegCloseKey(outkey);
	return false;
}

int APIENTRY _tWinMain(_In_ HINSTANCE hInstance,
                     _In_opt_ HINSTANCE hPrevInstance,
                     _In_ LPTSTR    lpCmdLine,
                     _In_ int       nCmdShow)
{
	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);
	DWORD msiResult;

	// Parse command line
	// Acceptable options:
	// /TEST - test managementagent(x86/x64).msi based install
	// /quiet - silent, noninteractive install
	// /passive - passive install
	// /norestart - does not restart vm
	// /forcerestart - will always restart vm

	arguments args;

	BOOL installing = Installing();

	if (!parseCommandLine(&args))
		return 0;

	// We require .net to run

	if (!checkDotNet()) {
		ErrMsg(getBrandingString(BRANDING_errDotNetNeeded));
		return 0;
	}

	//Locate Files

	if (!getFileLocations()) {
		return 0;
	}

	msiResult = installMsi(&args);

	if (msiResult != ERROR_SUCCESS &&
			msiResult != ERROR_SUCCESS_REBOOT_INITIATED &&
			msiResult != ERROR_SUCCESS_REBOOT_REQUIRED) {
		ErrMsg(
			getBrandingString(BRANDING_errMSIInstallFail),
			msiResult,
			logfile
		);
		return msiResult;
	}

	if (installing) {
		waitForStatus();
	}

	return 0;
}

<|endoftext|>


<commit_before>/*
 * Copyright (c) 2004 The Regents of The University of Michigan
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions 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;
 * neither the name of the copyright holders nor the names of its
 * contributors may 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 COPYRIGHT
 * OWNER 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.
 *
 * Authors: Ali Saidi
 *          Nathan Binkert
 */

//The purpose of this file is to provide endainness conversion utility
//functions. Depending on the endianness of the guest system, either
//the LittleEndianGuest or BigEndianGuest namespace is used.

#ifndef __SIM_BYTE_SWAP_HH__
#define __SIM_BYTE_SWAP_HH__

#include "sim/host.hh"

// This lets us figure out what the byte order of the host system is
#if defined(linux)
#include <endian.h>
// If this is a linux system, lets used the optimized definitions if they exist.
// If one doesn't exist, we pretty much get what is listed below, so it all
// works out
#include <byteswap.h>
#else
#include <machine/endian.h>
#endif

#if defined(__APPLE__)
#include <libkern/OSByteOrder.h>
#endif

//These functions actually perform the swapping for parameters
//of various bit lengths
static inline uint64_t
swap_byte64(uint64_t x)
{
#if defined(linux)
    return bswap_64(x);
#elif defined(__APPLE__)
    return OSSwapInt64(x);
#else
    return  (uint64_t)((((uint64_t)(x) & 0xff) << 56) |
            ((uint64_t)(x) & 0xff00ULL) << 40 |
            ((uint64_t)(x) & 0xff0000ULL) << 24 |
            ((uint64_t)(x) & 0xff000000ULL) << 8 |
            ((uint64_t)(x) & 0xff00000000ULL) >> 8 |
            ((uint64_t)(x) & 0xff0000000000ULL) >> 24 |
            ((uint64_t)(x) & 0xff000000000000ULL) >> 40 |
            ((uint64_t)(x) & 0xff00000000000000ULL) >> 56) ;
#endif
}

static inline uint32_t
swap_byte32(uint32_t x)
{
#if defined(linux)
    return bswap_32(x);
#elif defined(__APPLE__)
    return OSSwapInt32(x);
#else
    return  (uint32_t)(((uint32_t)(x) & 0xff) << 24 |
            ((uint32_t)(x) & 0xff00) << 8 | ((uint32_t)(x) & 0xff0000) >> 8 |
            ((uint32_t)(x) & 0xff000000) >> 24);
#endif
}

static inline uint16_t
swap_byte16(uint16_t x)
{
#if defined(linux)
    return bswap_16(x);
#elif defined(__APPLE__)
    return OSSwapInt16(x);
#else
    return (uint16_t)(((uint16_t)(x) & 0xff) << 8 |
                      ((uint16_t)(x) & 0xff00) >> 8);
#endif
}

// This function lets the compiler figure out how to call the
// swap_byte functions above for different data types.  Since the
// sizeof() values are known at compiel time, it should inline to a
// direct call to the right swap_byteNN() function.
template <typename T>
static inline T swap_byte(T x) {
    if (sizeof(T) == 8)
        return swap_byte64((uint64_t)x);
    else if (sizeof(T) == 4)
        return swap_byte32((uint32_t)x);
    else if (sizeof(T) == 2)
        return swap_byte16((uint16_t)x);
    else if (sizeof(T) == 1)
        return x;
    else
        panic("Can't byte-swap values larger than 64 bits");
}

//The conversion functions with fixed endianness on both ends don't need to
//be in a namespace
template <typename T> static inline T betole(T value) {return swap_byte(value);}
template <typename T> static inline T letobe(T value) {return swap_byte(value);}

//For conversions not involving the guest system, we can define the functions
//conditionally based on the BYTE_ORDER macro and outside of the namespaces
#if BYTE_ORDER == BIG_ENDIAN
template <typename T> static inline T htole(T value) {return swap_byte(value);}
template <typename T> static inline T letoh(T value) {return swap_byte(value);}
template <typename T> static inline T htobe(T value) {return value;}
template <typename T> static inline T betoh(T value) {return value;}
#elif BYTE_ORDER == LITTLE_ENDIAN
template <typename T> static inline T htole(T value) {return value;}
template <typename T> static inline T letoh(T value) {return value;}
template <typename T> static inline T htobe(T value) {return swap_byte(value);}
template <typename T> static inline T betoh(T value) {return swap_byte(value);}
#else
        #error Invalid Endianess
#endif

namespace BigEndianGuest
{
        template <typename T>
        static inline T gtole(T value) {return betole(value);}
        template <typename T>
        static inline T letog(T value) {return letobe(value);}
        template <typename T>
        static inline T gtobe(T value) {return value;}
        template <typename T>
        static inline T betog(T value) {return value;}
        template <typename T>
        static inline T htog(T value) {return htobe(value);}
        template <typename T>
        static inline T gtoh(T value) {return betoh(value);}
}

namespace LittleEndianGuest
{
        template <typename T>
        static inline T gtole(T value) {return value;}
        template <typename T>
        static inline T letog(T value) {return value;}
        template <typename T>
        static inline T gtobe(T value) {return letobe(value);}
        template <typename T>
        static inline T betog(T value) {return betole(value);}
        template <typename T>
        static inline T htog(T value) {return htole(value);}
        template <typename T>
        static inline T gtoh(T value) {return letoh(value);}
}
#endif // __SIM_BYTE_SWAP_HH__
<commit_msg>include misc.hh for panic<commit_after>/*
 * Copyright (c) 2004 The Regents of The University of Michigan
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions 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;
 * neither the name of the copyright holders nor the names of its
 * contributors may 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 COPYRIGHT
 * OWNER 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.
 *
 * Authors: Ali Saidi
 *          Nathan Binkert
 */

//The purpose of this file is to provide endainness conversion utility
//functions. Depending on the endianness of the guest system, either
//the LittleEndianGuest or BigEndianGuest namespace is used.

#ifndef __SIM_BYTE_SWAP_HH__
#define __SIM_BYTE_SWAP_HH__

#include "base/misc.hh"
#include "sim/host.hh"

// This lets us figure out what the byte order of the host system is
#if defined(linux)
#include <endian.h>
// If this is a linux system, lets used the optimized definitions if they exist.
// If one doesn't exist, we pretty much get what is listed below, so it all
// works out
#include <byteswap.h>
#else
#include <machine/endian.h>
#endif

#if defined(__APPLE__)
#include <libkern/OSByteOrder.h>
#endif

//These functions actually perform the swapping for parameters
//of various bit lengths
static inline uint64_t
swap_byte64(uint64_t x)
{
#if defined(linux)
    return bswap_64(x);
#elif defined(__APPLE__)
    return OSSwapInt64(x);
#else
    return  (uint64_t)((((uint64_t)(x) & 0xff) << 56) |
            ((uint64_t)(x) & 0xff00ULL) << 40 |
            ((uint64_t)(x) & 0xff0000ULL) << 24 |
            ((uint64_t)(x) & 0xff000000ULL) << 8 |
            ((uint64_t)(x) & 0xff00000000ULL) >> 8 |
            ((uint64_t)(x) & 0xff0000000000ULL) >> 24 |
            ((uint64_t)(x) & 0xff000000000000ULL) >> 40 |
            ((uint64_t)(x) & 0xff00000000000000ULL) >> 56) ;
#endif
}

static inline uint32_t
swap_byte32(uint32_t x)
{
#if defined(linux)
    return bswap_32(x);
#elif defined(__APPLE__)
    return OSSwapInt32(x);
#else
    return  (uint32_t)(((uint32_t)(x) & 0xff) << 24 |
            ((uint32_t)(x) & 0xff00) << 8 | ((uint32_t)(x) & 0xff0000) >> 8 |
            ((uint32_t)(x) & 0xff000000) >> 24);
#endif
}

static inline uint16_t
swap_byte16(uint16_t x)
{
#if defined(linux)
    return bswap_16(x);
#elif defined(__APPLE__)
    return OSSwapInt16(x);
#else
    return (uint16_t)(((uint16_t)(x) & 0xff) << 8 |
                      ((uint16_t)(x) & 0xff00) >> 8);
#endif
}

// This function lets the compiler figure out how to call the
// swap_byte functions above for different data types.  Since the
// sizeof() values are known at compiel time, it should inline to a
// direct call to the right swap_byteNN() function.
template <typename T>
static inline T swap_byte(T x) {
    if (sizeof(T) == 8)
        return swap_byte64((uint64_t)x);
    else if (sizeof(T) == 4)
        return swap_byte32((uint32_t)x);
    else if (sizeof(T) == 2)
        return swap_byte16((uint16_t)x);
    else if (sizeof(T) == 1)
        return x;
    else
        panic("Can't byte-swap values larger than 64 bits");
}

//The conversion functions with fixed endianness on both ends don't need to
//be in a namespace
template <typename T> static inline T betole(T value) {return swap_byte(value);}
template <typename T> static inline T letobe(T value) {return swap_byte(value);}

//For conversions not involving the guest system, we can define the functions
//conditionally based on the BYTE_ORDER macro and outside of the namespaces
#if BYTE_ORDER == BIG_ENDIAN
template <typename T> static inline T htole(T value) {return swap_byte(value);}
template <typename T> static inline T letoh(T value) {return swap_byte(value);}
template <typename T> static inline T htobe(T value) {return value;}
template <typename T> static inline T betoh(T value) {return value;}
#elif BYTE_ORDER == LITTLE_ENDIAN
template <typename T> static inline T htole(T value) {return value;}
template <typename T> static inline T letoh(T value) {return value;}
template <typename T> static inline T htobe(T value) {return swap_byte(value);}
template <typename T> static inline T betoh(T value) {return swap_byte(value);}
#else
        #error Invalid Endianess
#endif

namespace BigEndianGuest
{
        template <typename T>
        static inline T gtole(T value) {return betole(value);}
        template <typename T>
        static inline T letog(T value) {return letobe(value);}
        template <typename T>
        static inline T gtobe(T value) {return value;}
        template <typename T>
        static inline T betog(T value) {return value;}
        template <typename T>
        static inline T htog(T value) {return htobe(value);}
        template <typename T>
        static inline T gtoh(T value) {return betoh(value);}
}

namespace LittleEndianGuest
{
        template <typename T>
        static inline T gtole(T value) {return value;}
        template <typename T>
        static inline T letog(T value) {return value;}
        template <typename T>
        static inline T gtobe(T value) {return letobe(value);}
        template <typename T>
        static inline T betog(T value) {return betole(value);}
        template <typename T>
        static inline T htog(T value) {return htole(value);}
        template <typename T>
        static inline T gtoh(T value) {return letoh(value);}
}
#endif // __SIM_BYTE_SWAP_HH__
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

#ifndef __SLAVE_FLAGS_HPP__
#define __SLAVE_FLAGS_HPP__

#include <string>

#include <stout/duration.hpp>
#include <stout/flags.hpp>
#include <stout/option.hpp>

#include "logging/flags.hpp"

#include "slave/constants.hpp"

namespace mesos {
namespace internal {
namespace slave {

class Flags : public logging::Flags
{
public:
  Flags()
  {
    add(&Flags::hostname,
        "hostname",
        "The hostname the slave should report.\n"
        "If left unset, system hostname will be used (recommended).");

    add(&Flags::version,
        "version",
        "Show version and exit.",
        false);

    // TODO(benh): Is there a way to specify units for the resources?
    add(&Flags::resources,
        "resources",
        "Total consumable resources per slave, in\n"
        "the form 'name(role):value;name(role):value...'.");

    add(&Flags::isolation,
        "isolation",
        "Isolation mechanisms to use, e.g., 'posix/cpu,posix/mem', or\n"
        "'cgroups/cpu,cgroups/mem', or network/port_mapping\n"
        "(configure with flag: --with-network-isolator to enable),\n"
        "or 'external'.",
        "posix/cpu,posix/mem");

    add(&Flags::default_role,
        "default_role",
        "Any resources in the --resources flag that\n"
        "omit a role, as well as any resources that\n"
        "are not present in --resources but that are\n"
        "automatically detected, will be assigned to\n"
        "this role.",
        "*");

    add(&Flags::attributes,
        "attributes",
        "Attributes of machine, in the form:\n"
        "rack:2 or 'rack:2,u:1'");

    add(&Flags::work_dir,
        "work_dir",
        "Directory path to place framework work directories\n",
        "/tmp/mesos");

    add(&Flags::launcher_dir, // TODO(benh): This needs a better name.
        "launcher_dir",
        "Directory path of Mesos binaries",
        PKGLIBEXECDIR);

    add(&Flags::hadoop_home,
        "hadoop_home",
        "Path to find Hadoop installed (for\n"
        "fetching framework executors from HDFS)\n"
        "(no default, look for HADOOP_HOME in\n"
        "environment or find hadoop on PATH)",
        "");

    add(&Flags::switch_user,
        "switch_user",
        "Whether to run tasks as the user who\n"
        "submitted them rather than the user running\n"
        "the slave (requires setuid permission)",
        true);

    add(&Flags::frameworks_home,
        "frameworks_home",
        "Directory path prepended to relative executor URIs",
        "");

    add(&Flags::registration_backoff_factor,
        "registration_backoff_factor",
        "Slave initially picks a random amount of time between [0, b], where\n"
        "b = registration_backoff_factor, to (re-)register with a new master.\n"
        "Subsequent retries are exponentially backed off based on this\n"
        "interval (e.g., 1st retry uses a random value between [0, b * 2^1],\n"
        "2nd retry between [0, b * 2^2], 3rd retry between [0, b * 2^3] etc)\n"
        "up to a maximum of " + stringify(REGISTER_RETRY_INTERVAL_MAX),
        REGISTRATION_BACKOFF_FACTOR);

    add(&Flags::executor_registration_timeout,
        "executor_registration_timeout",
        "Amount of time to wait for an executor\n"
        "to register with the slave before considering it hung and\n"
        "shutting it down (e.g., 60secs, 3mins, etc)",
        EXECUTOR_REGISTRATION_TIMEOUT);

    add(&Flags::executor_shutdown_grace_period,
        "executor_shutdown_grace_period",
        "Amount of time to wait for an executor\n"
        "to shut down (e.g., 60secs, 3mins, etc)",
        EXECUTOR_SHUTDOWN_GRACE_PERIOD);

    add(&Flags::gc_delay,
        "gc_delay",
        "Maximum amount of time to wait before cleaning up\n"
        "executor directories (e.g., 3days, 2weeks, etc).\n"
        "Note that this delay may be shorter depending on\n"
        "the available disk usage.",
        GC_DELAY);

    add(&Flags::disk_watch_interval,
        "disk_watch_interval",
        "Periodic time interval (e.g., 10secs, 2mins, etc)\n"
        "to check the disk usage",
        DISK_WATCH_INTERVAL);

    add(&Flags::resource_monitoring_interval,
        "resource_monitoring_interval",
        "Periodic time interval for monitoring executor\n"
        "resource usage (e.g., 10secs, 1min, etc)",
        RESOURCE_MONITORING_INTERVAL);

    // TODO(vinod): Consider killing this flag and always checkpoint.
    add(&Flags::checkpoint,
        "checkpoint",
        "Whether to checkpoint slave and frameworks information\n"
        "to disk. This enables a restarted slave to recover\n"
        "status updates and reconnect with (--recover=reconnect) or\n"
        "kill (--recover=cleanup) old executors",
        true);

    add(&Flags::recover,
        "recover",
        "Whether to recover status updates and reconnect with old executors.\n"
        "Valid values for 'recover' are\n"
        "reconnect: Reconnect with any old live executors.\n"
        "cleanup  : Kill any old live executors and exit.\n"
        "           Use this option when doing an incompatible slave\n"
        "           or executor upgrade!).\n"
        "NOTE: If checkpointed slave doesn't exist, no recovery is performed\n"
        "      and the slave registers with the master as a new slave.",
        "reconnect");

    add(&Flags::recovery_timeout,
        "recovery_timeout",
        "Amount of time alloted for the slave to recover. If the slave takes\n"
        "longer than recovery_timeout to recover, any executors that are\n"
        "waiting to reconnect to the slave will self-terminate.\n"
        "NOTE: This flag is only applicable when checkpoint is enabled.\n",
        RECOVERY_TIMEOUT);

    add(&Flags::strict,
        "strict",
        "If strict=true, any and all recovery errors are considered fatal.\n"
        "If strict=false, any expected errors (e.g., slave cannot recover\n"
        "information about an executor, because the slave died right before\n"
        "the executor registered.) during recovery are ignored and as much\n"
        "state as possible is recovered.\n",
        true);

#ifdef __linux__
    add(&Flags::cgroups_hierarchy,
        "cgroups_hierarchy",
        "The path to the cgroups hierarchy root\n",
        "/sys/fs/cgroup");

    add(&Flags::cgroups_root,
        "cgroups_root",
        "Name of the root cgroup\n",
        "mesos");

    add(&Flags::cgroups_subsystems,
        "cgroups_subsystems",
        "This flag has been deprecated and is no longer used,\n"
        "please update your flags");

    add(&Flags::cgroups_enable_cfs,
        "cgroups_enable_cfs",
        "Cgroups feature flag to enable hard limits on CPU resources\n"
        "via the CFS bandwidth limiting subfeature.\n",
        false);

    // TODO(antonl): Set default to true in future releases.
    add(&Flags::cgroups_limit_swap,
        "cgroups_limit_swap",
        "Cgroups feature flag to enable memory limits on both memory and\n"
        "swap instead of just memory.\n",
        false);

    add(&Flags::slave_subsystems,
        "slave_subsystems",
        "List of comma-separated cgroup subsystems to run the slave binary\n"
        "in, e.g., 'memory,cpuacct'. The default is none.\n"
        "Present functionality is intended for resource monitoring and\n"
        "no cgroup limits are set, they are inherited from the root mesos\n"
        "cgroup.");

    add(&Flags::perf_events,
        "perf_events",
        "List of command-separated perf events to sample for each container\n"
        "when using the perf_event isolator. Default is none.\n"
        "Run command 'perf list' to see all events. Event names are\n"
        "sanitized by downcasing and replacing hyphens with underscores\n"
        "when reported in the PerfStatistics protobuf, e.g., cpu-cycles\n"
        "becomes cpu_cycles; see the PerfStatistics protobuf for all names.");

    add(&Flags::perf_interval,
        "perf_interval",
        "Interval between the start of perf stat samples. Perf samples are\n"
        "obtained periodically according to perf_interval and the most\n"
        "recently obtained sample is returned rather than sampling on\n"
        "demand. For this reason, perf_interval is independent of the\n"
        "resource monitoring interval",
        Seconds(60));

    add(&Flags::perf_duration,
        "perf_duration",
        "Duration of a perf stat sample. The duration must be less\n"
        "that the perf_interval.",
        Seconds(10));
#endif

    add(&Flags::credential,
        "credential",
        "Either a path to a text with a single line\n"
        "containing 'principal' and 'secret' separated by "
        "whitespace.\n"
        "Or a path containing the JSON "
        "formatted information used for one credential.\n"
        "Path could be of the form 'file:///path/to/file' or '/path/to/file'."
        "\n"
        "Example:\n"
        "{\n"
        "    \"principal\": \"username\",\n"
        "    \"secret\": \"secret\",\n"
        "}");

    add(&Flags::containerizer_path,
        "containerizer_path",
        "The path to the external containerizer executable used when\n"
        "external isolation is activated (--isolation=external).\n");

    add(&Flags::containerizers,
        "containerizers",
        "Comma separated list of containerizer implementations\n"
        "to compose in order to provide containerization.\n"
        "Available options are 'mesos', 'external', and\n"
        "'docker' (on Linux). The order the containerizers\n"
        "are specified is the order they are tried\n"
        "(--containerizers=mesos).\n",
        "mesos");

    add(&Flags::default_container_image,
        "default_container_image",
        "The default container image to use if not specified by a task,\n"
        "when using external containerizer");

    add(&Flags::docker,
        "docker",
        "The absolute path to the docker executable for docker containerizer.",
        "docker");

    add(&Flags::docker_sandbox_directory,
        "docker_sandbox_directory",
        "The absolute path for the directory in the container where the sandbox\n"
        "is mapped to",
        "/mnt/mesos/sandbox");

#ifdef WITH_NETWORK_ISOLATOR
    add(&Flags::ephemeral_ports_per_container,
        "ephemeral_ports_per_container",
        "Number of ephemeral ports allocated to a container by the network\n"
        "isolator. This number has to be a power of 2.\n",
        DEFAULT_EPHEMERAL_PORTS_PER_CONTAINER);

    add(&Flags::eth0_name,
        "eth0_name",
        "The name of the public network interface (e.g., eth0). If it is\n"
        "not specified, the network isolator will try to guess it based\n"
        "on the host default gateway.");

    add(&Flags::lo_name,
        "lo_name",
        "The name of the loopback network interface (e.g., lo). If it is\n"
        "not specified, the network isolator will try to guess it.");
#endif // WITH_NETWORK_ISOLATOR
  }

  bool version;
  Option<std::string> hostname;
  Option<std::string> resources;
  std::string isolation;
  std::string default_role;
  Option<std::string> attributes;
  std::string work_dir;
  std::string launcher_dir;
  std::string hadoop_home; // TODO(benh): Make an Option.
  bool switch_user;
  std::string frameworks_home;  // TODO(benh): Make an Option.
  Duration registration_backoff_factor;
  Duration executor_registration_timeout;
  Duration executor_shutdown_grace_period;
  Duration gc_delay;
  Duration disk_watch_interval;
  Duration resource_monitoring_interval;
  bool checkpoint;
  std::string recover;
  Duration recovery_timeout;
  bool strict;
  Duration register_retry_interval_min;
#ifdef __linux__
  std::string cgroups_hierarchy;
  std::string cgroups_root;
  Option<std::string> cgroups_subsystems;
  bool cgroups_enable_cfs;
  bool cgroups_limit_swap;
  Option<std::string> slave_subsystems;
  Option<std::string> perf_events;
  Duration perf_interval;
  Duration perf_duration;
#endif
  Option<std::string> credential;
  Option<std::string> containerizer_path;
  std::string containerizers;
  Option<std::string> default_container_image;
  std::string docker;
  std::string docker_sandbox_directory;
#ifdef WITH_NETWORK_ISOLATOR
  uint16_t ephemeral_ports_per_container;
  Option<std::string> eth0_name;
  Option<std::string> lo_name;
#endif
};

} // namespace slave {
} // namespace internal {
} // namespace mesos {

#endif // __SLAVE_FLAGS_HPP__
<commit_msg>Fixed style issues in slave/flags.hpp.<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

#ifndef __SLAVE_FLAGS_HPP__
#define __SLAVE_FLAGS_HPP__

#include <string>

#include <stout/duration.hpp>
#include <stout/flags.hpp>
#include <stout/option.hpp>

#include "logging/flags.hpp"

#include "slave/constants.hpp"

namespace mesos {
namespace internal {
namespace slave {

class Flags : public logging::Flags
{
public:
  Flags()
  {
    add(&Flags::hostname,
        "hostname",
        "The hostname the slave should report.\n"
        "If left unset, system hostname will be used (recommended).");

    add(&Flags::version,
        "version",
        "Show version and exit.",
        false);

    // TODO(benh): Is there a way to specify units for the resources?
    add(&Flags::resources,
        "resources",
        "Total consumable resources per slave, in\n"
        "the form 'name(role):value;name(role):value...'.");

    add(&Flags::isolation,
        "isolation",
        "Isolation mechanisms to use, e.g., 'posix/cpu,posix/mem', or\n"
        "'cgroups/cpu,cgroups/mem', or network/port_mapping\n"
        "(configure with flag: --with-network-isolator to enable),\n"
        "or 'external'.",
        "posix/cpu,posix/mem");

    add(&Flags::default_role,
        "default_role",
        "Any resources in the --resources flag that\n"
        "omit a role, as well as any resources that\n"
        "are not present in --resources but that are\n"
        "automatically detected, will be assigned to\n"
        "this role.",
        "*");

    add(&Flags::attributes,
        "attributes",
        "Attributes of machine, in the form:\n"
        "rack:2 or 'rack:2,u:1'");

    add(&Flags::work_dir,
        "work_dir",
        "Directory path to place framework work directories\n",
        "/tmp/mesos");

    add(&Flags::launcher_dir, // TODO(benh): This needs a better name.
        "launcher_dir",
        "Directory path of Mesos binaries",
        PKGLIBEXECDIR);

    add(&Flags::hadoop_home,
        "hadoop_home",
        "Path to find Hadoop installed (for\n"
        "fetching framework executors from HDFS)\n"
        "(no default, look for HADOOP_HOME in\n"
        "environment or find hadoop on PATH)",
        "");

    add(&Flags::switch_user,
        "switch_user",
        "Whether to run tasks as the user who\n"
        "submitted them rather than the user running\n"
        "the slave (requires setuid permission)",
        true);

    add(&Flags::frameworks_home,
        "frameworks_home",
        "Directory path prepended to relative executor URIs",
        "");

    add(&Flags::registration_backoff_factor,
        "registration_backoff_factor",
        "Slave initially picks a random amount of time between [0, b], where\n"
        "b = registration_backoff_factor, to (re-)register with a new master.\n"
        "Subsequent retries are exponentially backed off based on this\n"
        "interval (e.g., 1st retry uses a random value between [0, b * 2^1],\n"
        "2nd retry between [0, b * 2^2], 3rd retry between [0, b * 2^3] etc)\n"
        "up to a maximum of " + stringify(REGISTER_RETRY_INTERVAL_MAX),
        REGISTRATION_BACKOFF_FACTOR);

    add(&Flags::executor_registration_timeout,
        "executor_registration_timeout",
        "Amount of time to wait for an executor\n"
        "to register with the slave before considering it hung and\n"
        "shutting it down (e.g., 60secs, 3mins, etc)",
        EXECUTOR_REGISTRATION_TIMEOUT);

    add(&Flags::executor_shutdown_grace_period,
        "executor_shutdown_grace_period",
        "Amount of time to wait for an executor\n"
        "to shut down (e.g., 60secs, 3mins, etc)",
        EXECUTOR_SHUTDOWN_GRACE_PERIOD);

    add(&Flags::gc_delay,
        "gc_delay",
        "Maximum amount of time to wait before cleaning up\n"
        "executor directories (e.g., 3days, 2weeks, etc).\n"
        "Note that this delay may be shorter depending on\n"
        "the available disk usage.",
        GC_DELAY);

    add(&Flags::disk_watch_interval,
        "disk_watch_interval",
        "Periodic time interval (e.g., 10secs, 2mins, etc)\n"
        "to check the disk usage",
        DISK_WATCH_INTERVAL);

    add(&Flags::resource_monitoring_interval,
        "resource_monitoring_interval",
        "Periodic time interval for monitoring executor\n"
        "resource usage (e.g., 10secs, 1min, etc)",
        RESOURCE_MONITORING_INTERVAL);

    // TODO(vinod): Consider killing this flag and always checkpoint.
    add(&Flags::checkpoint,
        "checkpoint",
        "Whether to checkpoint slave and frameworks information\n"
        "to disk. This enables a restarted slave to recover\n"
        "status updates and reconnect with (--recover=reconnect) or\n"
        "kill (--recover=cleanup) old executors",
        true);

    add(&Flags::recover,
        "recover",
        "Whether to recover status updates and reconnect with old executors.\n"
        "Valid values for 'recover' are\n"
        "reconnect: Reconnect with any old live executors.\n"
        "cleanup  : Kill any old live executors and exit.\n"
        "           Use this option when doing an incompatible slave\n"
        "           or executor upgrade!).\n"
        "NOTE: If checkpointed slave doesn't exist, no recovery is performed\n"
        "      and the slave registers with the master as a new slave.",
        "reconnect");

    add(&Flags::recovery_timeout,
        "recovery_timeout",
        "Amount of time alloted for the slave to recover. If the slave takes\n"
        "longer than recovery_timeout to recover, any executors that are\n"
        "waiting to reconnect to the slave will self-terminate.\n"
        "NOTE: This flag is only applicable when checkpoint is enabled.\n",
        RECOVERY_TIMEOUT);

    add(&Flags::strict,
        "strict",
        "If strict=true, any and all recovery errors are considered fatal.\n"
        "If strict=false, any expected errors (e.g., slave cannot recover\n"
        "information about an executor, because the slave died right before\n"
        "the executor registered.) during recovery are ignored and as much\n"
        "state as possible is recovered.\n",
        true);

#ifdef __linux__
    add(&Flags::cgroups_hierarchy,
        "cgroups_hierarchy",
        "The path to the cgroups hierarchy root\n",
        "/sys/fs/cgroup");

    add(&Flags::cgroups_root,
        "cgroups_root",
        "Name of the root cgroup\n",
        "mesos");

    add(&Flags::cgroups_subsystems,
        "cgroups_subsystems",
        "This flag has been deprecated and is no longer used,\n"
        "please update your flags");

    add(&Flags::cgroups_enable_cfs,
        "cgroups_enable_cfs",
        "Cgroups feature flag to enable hard limits on CPU resources\n"
        "via the CFS bandwidth limiting subfeature.\n",
        false);

    // TODO(antonl): Set default to true in future releases.
    add(&Flags::cgroups_limit_swap,
        "cgroups_limit_swap",
        "Cgroups feature flag to enable memory limits on both memory and\n"
        "swap instead of just memory.\n",
        false);

    add(&Flags::slave_subsystems,
        "slave_subsystems",
        "List of comma-separated cgroup subsystems to run the slave binary\n"
        "in, e.g., 'memory,cpuacct'. The default is none.\n"
        "Present functionality is intended for resource monitoring and\n"
        "no cgroup limits are set, they are inherited from the root mesos\n"
        "cgroup.");

    add(&Flags::perf_events,
        "perf_events",
        "List of command-separated perf events to sample for each container\n"
        "when using the perf_event isolator. Default is none.\n"
        "Run command 'perf list' to see all events. Event names are\n"
        "sanitized by downcasing and replacing hyphens with underscores\n"
        "when reported in the PerfStatistics protobuf, e.g., cpu-cycles\n"
        "becomes cpu_cycles; see the PerfStatistics protobuf for all names.");

    add(&Flags::perf_interval,
        "perf_interval",
        "Interval between the start of perf stat samples. Perf samples are\n"
        "obtained periodically according to perf_interval and the most\n"
        "recently obtained sample is returned rather than sampling on\n"
        "demand. For this reason, perf_interval is independent of the\n"
        "resource monitoring interval",
        Seconds(60));

    add(&Flags::perf_duration,
        "perf_duration",
        "Duration of a perf stat sample. The duration must be less\n"
        "that the perf_interval.",
        Seconds(10));
#endif

    add(&Flags::credential,
        "credential",
        "Either a path to a text with a single line\n"
        "containing 'principal' and 'secret' separated by "
        "whitespace.\n"
        "Or a path containing the JSON "
        "formatted information used for one credential.\n"
        "Path could be of the form 'file:///path/to/file' or '/path/to/file'."
        "\n"
        "Example:\n"
        "{\n"
        "    \"principal\": \"username\",\n"
        "    \"secret\": \"secret\",\n"
        "}");

    add(&Flags::containerizer_path,
        "containerizer_path",
        "The path to the external containerizer executable used when\n"
        "external isolation is activated (--isolation=external).\n");

    add(&Flags::containerizers,
        "containerizers",
        "Comma separated list of containerizer implementations\n"
        "to compose in order to provide containerization.\n"
        "Available options are 'mesos', 'external', and\n"
        "'docker' (on Linux). The order the containerizers\n"
        "are specified is the order they are tried\n"
        "(--containerizers=mesos).\n",
        "mesos");

    add(&Flags::default_container_image,
        "default_container_image",
        "The default container image to use if not specified by a task,\n"
        "when using external containerizer.\n");

    add(&Flags::docker,
        "docker",
        "The absolute path to the docker executable for docker\n"
        "containerizer.\n",
        "docker");

    add(&Flags::docker_sandbox_directory,
        "docker_sandbox_directory",
        "The absolute path for the directory in the container where the\n"
        "sandbox is mapped to.\n",
        "/mnt/mesos/sandbox");

#ifdef WITH_NETWORK_ISOLATOR
    add(&Flags::ephemeral_ports_per_container,
        "ephemeral_ports_per_container",
        "Number of ephemeral ports allocated to a container by the network\n"
        "isolator. This number has to be a power of 2.\n",
        DEFAULT_EPHEMERAL_PORTS_PER_CONTAINER);

    add(&Flags::eth0_name,
        "eth0_name",
        "The name of the public network interface (e.g., eth0). If it is\n"
        "not specified, the network isolator will try to guess it based\n"
        "on the host default gateway.");

    add(&Flags::lo_name,
        "lo_name",
        "The name of the loopback network interface (e.g., lo). If it is\n"
        "not specified, the network isolator will try to guess it.");
#endif // WITH_NETWORK_ISOLATOR
  }

  bool version;
  Option<std::string> hostname;
  Option<std::string> resources;
  std::string isolation;
  std::string default_role;
  Option<std::string> attributes;
  std::string work_dir;
  std::string launcher_dir;
  std::string hadoop_home; // TODO(benh): Make an Option.
  bool switch_user;
  std::string frameworks_home;  // TODO(benh): Make an Option.
  Duration registration_backoff_factor;
  Duration executor_registration_timeout;
  Duration executor_shutdown_grace_period;
  Duration gc_delay;
  Duration disk_watch_interval;
  Duration resource_monitoring_interval;
  bool checkpoint;
  std::string recover;
  Duration recovery_timeout;
  bool strict;
  Duration register_retry_interval_min;
#ifdef __linux__
  std::string cgroups_hierarchy;
  std::string cgroups_root;
  Option<std::string> cgroups_subsystems;
  bool cgroups_enable_cfs;
  bool cgroups_limit_swap;
  Option<std::string> slave_subsystems;
  Option<std::string> perf_events;
  Duration perf_interval;
  Duration perf_duration;
#endif
  Option<std::string> credential;
  Option<std::string> containerizer_path;
  std::string containerizers;
  Option<std::string> default_container_image;
  std::string docker;
  std::string docker_sandbox_directory;
#ifdef WITH_NETWORK_ISOLATOR
  uint16_t ephemeral_ports_per_container;
  Option<std::string> eth0_name;
  Option<std::string> lo_name;
#endif
};

} // namespace slave {
} // namespace internal {
} // namespace mesos {

#endif // __SLAVE_FLAGS_HPP__
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

#ifndef __SLAVE_STATE_HPP__
#define __SLAVE_STATE_HPP__

#include <unistd.h>

#include <vector>

#include <mesos/resources.hpp>
#include <mesos/type_utils.hpp>

#include <process/pid.hpp>

#include <stout/foreach.hpp>
#include <stout/hashmap.hpp>
#include <stout/hashset.hpp>
#include <stout/os.hpp>
#include <stout/protobuf.hpp>
#include <stout/strings.hpp>
#include <stout/utils.hpp>
#include <stout/uuid.hpp>

#include "messages/messages.hpp"

namespace mesos {
namespace internal {
namespace slave {
namespace state {

// Forward declarations.
struct State;
struct SlaveState;
struct FrameworkState;
struct ExecutorState;
struct RunState;
struct TaskState;
struct ResourcesState;

// This function performs recovery from the state stored at 'rootDir'.
// If the 'strict' flag is set, any errors encountered while
// recovering a state are considered fatal and hence the recovery is
// short-circuited and returns an error. There might be orphaned
// executors that need to be manually cleaned up. If 'strict' flag is
// not set, any errors encountered are considered non-fatal and the
// recovery continues by recovering as much of the state as possible,
// while increasing the 'errors' count. Note that 'errors' on a struct
// includes the 'errors' encountered recursively. In other words,
// 'State.errors' is the sum total of all recovery errors. If the
// machine has rebooted since the last slave run, None is returned.
Result<State> recover(const std::string& rootDir, bool strict);

namespace internal {

inline Try<Nothing> checkpoint(
    const std::string& path,
    const std::string& message)
{
  return ::os::write(path, message);
}

inline Try<Nothing> checkpoint(
    const std::string& path,
    const google::protobuf::Message& message)
{
  return ::protobuf::write(path, message);
}

template <typename T>
Try<Nothing> checkpoint(
    const std::string& path,
    const google::protobuf::RepeatedPtrField<T>& messages)
{
  return ::protobuf::write(path, messages);
}

inline Try<Nothing> checkpoint(
    const std::string& path,
    const Resources& resources)
{
  const google::protobuf::RepeatedPtrField<Resource>& messages = resources;
  return checkpoint(path, messages);
}

}  // namespace internal {

// Thin wrapper to checkpoint data to disk and perform the necessary
// error checking. It checkpoints an instance of T at the given path.
// We can checkpoint anything as long as T is supported by
// internal::checkpoint. Currently the list of supported Ts are:
//   - std::string
//   - google::protobuf::Message
//   - google::protobuf::RepeatedPtrField<T>
//   - mesos::Resources
//
// NOTE: We provide atomic (all-or-nothing) semantics here by always
// writing to a temporary file first then use os::rename to atomically
// move it to the desired path.
template <typename T>
Try<Nothing> checkpoint(const std::string& path, const T& t)
{
  // Create the base directory.
  Try<std::string> base = os::dirname(path);
  if (base.isError()) {
    return Error("Failed to get the base directory path: " + base.error());
  }

  Try<Nothing> mkdir = os::mkdir(base.get());
  if (mkdir.isError()) {
    return Error("Failed to create directory '" + base.get() +
                 "': " + mkdir.error());
  }

  // NOTE: We create the temporary file at 'base/XXXXXX' to make sure
  // rename below does not cross devices (MESOS-2319).
  //
  // TODO(jieyu): It's possible that the temporary file becomes
  // dangling if slave crashes or restarts while checkpointing.
  // Consider adding a way to garbage collect them.
  Try<std::string> temp = os::mktemp(path::join(base.get(), "XXXXXX"));
  if (temp.isError()) {
    return Error("Failed to create temporary file: " + temp.error());
  }

  // Now checkpoint the instance of T to the temporary file.
  Try<Nothing> checkpoint = internal::checkpoint(temp.get(), t);
  if (checkpoint.isError()) {
    // Try removing the temporary file on error.
    os::rm(temp.get());

    return Error("Failed to write temporary file '" + temp.get() +
                 "': " + checkpoint.error());
  }

  // Rename the temporary file to the path.
  Try<Nothing> rename = os::rename(temp.get(), path);
  if (rename.isError()) {
    // Try removing the temporary file on error.
    os::rm(temp.get());

    return Error("Failed to rename '" + temp.get() + "' to '" +
                 path + "': " + rename.error());
  }

  return Nothing();
}

// The top level state. Each of the structs below (recursively)
// recover the checkpointed state.
struct State
{
  State() : errors(0) {}

  Option<ResourcesState> resources;
  Option<SlaveState> slave;

  // TODO(jieyu): Consider using a vector of Option<Error> here so
  // that we can print all the errors. This also applies to all the
  // State structs below.
  unsigned int errors;
};

struct ResourcesState
{
  ResourcesState() : errors(0) {}

  static Try<ResourcesState> recover(
      const std::string& rootDir,
      bool strict);

  Resources resources;
  unsigned int errors;
};

struct SlaveState
{
  SlaveState () : errors(0) {}

  static Try<SlaveState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      bool strict);

  SlaveID id;
  Option<SlaveInfo> info;
  hashmap<FrameworkID, FrameworkState> frameworks;
  unsigned int errors;
};

struct FrameworkState
{
  FrameworkState () : errors(0) {}

  static Try<FrameworkState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      bool strict);

  FrameworkID id;
  Option<FrameworkInfo> info;
  Option<process::UPID> pid;
  hashmap<ExecutorID, ExecutorState> executors;
  unsigned int errors;
};

struct ExecutorState
{
  ExecutorState () : errors(0) {}

  static Try<ExecutorState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      bool strict);

  ExecutorID id;
  Option<ExecutorInfo> info;
  Option<ContainerID> latest;
  hashmap<ContainerID, RunState> runs;
  unsigned int errors;
};

struct RunState
{
  RunState () : completed(false), errors(0) {}

  static Try<RunState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      const ContainerID& containerId,
      bool strict);

  Option<ContainerID> id;
  hashmap<TaskID, TaskState> tasks;
  Option<pid_t> forkedPid;
  Option<process::UPID> libprocessPid;

  // Executor terminated and all its updates acknowledged.
  bool completed;

  unsigned int errors;
};

struct TaskState
{
  TaskState () : errors(0) {}

  static Try<TaskState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      const ContainerID& containerId,
      const TaskID& taskId,
      bool strict);

  TaskID id;
  Option<Task> info;
  std::vector<StatusUpdate> updates;
  hashset<UUID> acks;
  unsigned int errors;
};

} // namespace state {
} // namespace slave {
} // namespace internal {
} // namespace mesos {

#endif // __SLAVE_STATE_HPP__
<commit_msg>Re-order structs to prevent forward declaration dependency.<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

#ifndef __SLAVE_STATE_HPP__
#define __SLAVE_STATE_HPP__

#include <unistd.h>

#include <vector>

#include <mesos/resources.hpp>
#include <mesos/type_utils.hpp>

#include <process/pid.hpp>

#include <stout/foreach.hpp>
#include <stout/hashmap.hpp>
#include <stout/hashset.hpp>
#include <stout/os.hpp>
#include <stout/protobuf.hpp>
#include <stout/strings.hpp>
#include <stout/utils.hpp>
#include <stout/uuid.hpp>

#include "messages/messages.hpp"

namespace mesos {
namespace internal {
namespace slave {
namespace state {

// Forward declarations.
struct State;
struct SlaveState;
struct ResourcesState;
struct FrameworkState;
struct ExecutorState;
struct RunState;
struct TaskState;

// This function performs recovery from the state stored at 'rootDir'.
// If the 'strict' flag is set, any errors encountered while
// recovering a state are considered fatal and hence the recovery is
// short-circuited and returns an error. There might be orphaned
// executors that need to be manually cleaned up. If 'strict' flag is
// not set, any errors encountered are considered non-fatal and the
// recovery continues by recovering as much of the state as possible,
// while increasing the 'errors' count. Note that 'errors' on a struct
// includes the 'errors' encountered recursively. In other words,
// 'State.errors' is the sum total of all recovery errors. If the
// machine has rebooted since the last slave run, None is returned.
Result<State> recover(const std::string& rootDir, bool strict);

namespace internal {

inline Try<Nothing> checkpoint(
    const std::string& path,
    const std::string& message)
{
  return ::os::write(path, message);
}

inline Try<Nothing> checkpoint(
    const std::string& path,
    const google::protobuf::Message& message)
{
  return ::protobuf::write(path, message);
}

template <typename T>
Try<Nothing> checkpoint(
    const std::string& path,
    const google::protobuf::RepeatedPtrField<T>& messages)
{
  return ::protobuf::write(path, messages);
}

inline Try<Nothing> checkpoint(
    const std::string& path,
    const Resources& resources)
{
  const google::protobuf::RepeatedPtrField<Resource>& messages = resources;
  return checkpoint(path, messages);
}

}  // namespace internal {

// Thin wrapper to checkpoint data to disk and perform the necessary
// error checking. It checkpoints an instance of T at the given path.
// We can checkpoint anything as long as T is supported by
// internal::checkpoint. Currently the list of supported Ts are:
//   - std::string
//   - google::protobuf::Message
//   - google::protobuf::RepeatedPtrField<T>
//   - mesos::Resources
//
// NOTE: We provide atomic (all-or-nothing) semantics here by always
// writing to a temporary file first then use os::rename to atomically
// move it to the desired path.
template <typename T>
Try<Nothing> checkpoint(const std::string& path, const T& t)
{
  // Create the base directory.
  Try<std::string> base = os::dirname(path);
  if (base.isError()) {
    return Error("Failed to get the base directory path: " + base.error());
  }

  Try<Nothing> mkdir = os::mkdir(base.get());
  if (mkdir.isError()) {
    return Error("Failed to create directory '" + base.get() +
                 "': " + mkdir.error());
  }

  // NOTE: We create the temporary file at 'base/XXXXXX' to make sure
  // rename below does not cross devices (MESOS-2319).
  //
  // TODO(jieyu): It's possible that the temporary file becomes
  // dangling if slave crashes or restarts while checkpointing.
  // Consider adding a way to garbage collect them.
  Try<std::string> temp = os::mktemp(path::join(base.get(), "XXXXXX"));
  if (temp.isError()) {
    return Error("Failed to create temporary file: " + temp.error());
  }

  // Now checkpoint the instance of T to the temporary file.
  Try<Nothing> checkpoint = internal::checkpoint(temp.get(), t);
  if (checkpoint.isError()) {
    // Try removing the temporary file on error.
    os::rm(temp.get());

    return Error("Failed to write temporary file '" + temp.get() +
                 "': " + checkpoint.error());
  }

  // Rename the temporary file to the path.
  Try<Nothing> rename = os::rename(temp.get(), path);
  if (rename.isError()) {
    // Try removing the temporary file on error.
    os::rm(temp.get());

    return Error("Failed to rename '" + temp.get() + "' to '" +
                 path + "': " + rename.error());
  }

  return Nothing();
}

// NOTE: The *State structs (e.g., TaskState, RunState, etc) are
// defined in reverse dependency order because many of them have
// Option<*State> dependencies which means we need them declared in
// their entirety in order to compile because things like
// Option<*State> need to know the final size of the types.

struct TaskState
{
  TaskState () : errors(0) {}

  static Try<TaskState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      const ContainerID& containerId,
      const TaskID& taskId,
      bool strict);

  TaskID id;
  Option<Task> info;
  std::vector<StatusUpdate> updates;
  hashset<UUID> acks;
  unsigned int errors;
};

struct RunState
{
  RunState () : completed(false), errors(0) {}

  static Try<RunState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      const ContainerID& containerId,
      bool strict);

  Option<ContainerID> id;
  hashmap<TaskID, TaskState> tasks;
  Option<pid_t> forkedPid;
  Option<process::UPID> libprocessPid;

  // Executor terminated and all its updates acknowledged.
  bool completed;

  unsigned int errors;
};

struct ExecutorState
{
  ExecutorState () : errors(0) {}

  static Try<ExecutorState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      const ExecutorID& executorId,
      bool strict);

  ExecutorID id;
  Option<ExecutorInfo> info;
  Option<ContainerID> latest;
  hashmap<ContainerID, RunState> runs;
  unsigned int errors;
};

struct FrameworkState
{
  FrameworkState () : errors(0) {}

  static Try<FrameworkState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      const FrameworkID& frameworkId,
      bool strict);

  FrameworkID id;
  Option<FrameworkInfo> info;
  Option<process::UPID> pid;
  hashmap<ExecutorID, ExecutorState> executors;
  unsigned int errors;
};

struct ResourcesState
{
  ResourcesState() : errors(0) {}

  static Try<ResourcesState> recover(
      const std::string& rootDir,
      bool strict);

  Resources resources;
  unsigned int errors;
};

struct SlaveState
{
  SlaveState () : errors(0) {}

  static Try<SlaveState> recover(
      const std::string& rootDir,
      const SlaveID& slaveId,
      bool strict);

  SlaveID id;
  Option<SlaveInfo> info;
  hashmap<FrameworkID, FrameworkState> frameworks;
  unsigned int errors;
};

// The top level state. The members are child nodes in the tree. Each
// child node (recursively) recovers the checkpointed state.
struct State
{
  State() : errors(0) {}

  Option<ResourcesState> resources;
  Option<SlaveState> slave;

  // TODO(jieyu): Consider using a vector of Option<Error> here so
  // that we can print all the errors. This also applies to all the
  // State structs below.
  unsigned int errors;
};

} // namespace state {
} // namespace slave {
} // namespace internal {
} // namespace mesos {

#endif // __SLAVE_STATE_HPP__
<|endoftext|>


<commit_before>#include <stout/check.hpp>
#include <stout/foreach.hpp>
#include <stout/os.hpp>
#include <stout/path.hpp>
#include <stout/result.hpp>
#include <stout/uuid.hpp>

#ifdef __linux__
#include "linux/cgroups.hpp"
#endif

#include "slave/containerizer/containerizer.hpp"
#include "slave/containerizer/mesos_containerizer.hpp"

#include "tests/containerizer.hpp"
#include "tests/environment.hpp"
#include "tests/flags.hpp"
#include "tests/mesos.hpp"

using namespace process;

using testing::_;

namespace mesos {
namespace internal {
namespace tests {

MesosTest::MesosTest(const Option<zookeeper::URL>& url) : cluster(url) {}

master::Flags MesosTest::CreateMasterFlags()
{
  master::Flags flags;

  // Enable authentication.
  flags.authenticate = true;

  // Create a temporary work directory (removed by Environment).
  const Try<std::string>& directory = environment->mkdtemp();
  CHECK_SOME(directory) << "Failed to create temporary directory";

  flags.work_dir = directory.get();

  // Create a default credentials file.
  const std::string& path = path::join(directory.get(), "credentials");

  const std::string& credentials =
    DEFAULT_CREDENTIAL.principal() + " " + DEFAULT_CREDENTIAL.secret();

  CHECK_SOME(os::write(path, credentials))
    << "Failed to write credentials to '" << path << "'";

  flags.credentials = "file://" + path;

  return flags;
}

slave::Flags MesosTest::CreateSlaveFlags()
{
  slave::Flags flags;

  // Create a temporary work directory (removed by Environment).
  Try<std::string> directory = environment->mkdtemp();
  CHECK_SOME(directory) << "Failed to create temporary directory";

  flags.work_dir = directory.get();

  flags.launcher_dir = path::join(tests::flags.build_dir, "src");

  // TODO(vinod): Consider making this true and fixing the tests.
  flags.checkpoint = false;

  flags.resources = Option<std::string>(
      "cpus:2;mem:1024;disk:1024;ports:[31000-32000]");

  return flags;
}

Try<process::PID<master::Master> > MesosTest::StartMaster(
    const Option<master::Flags>& flags,
    bool wait)
{
  Future<Nothing> detected = FUTURE_DISPATCH(_, &master::Master::detected);

  Try<process::PID<master::Master> > master = cluster.masters.start(
      flags.isNone() ? CreateMasterFlags() : flags.get());

  // Wait until the leader is detected because otherwise this master
  // may reject authentication requests because it doesn't know it's
  // the leader yet [MESOS-881].
  if (wait && master.isSome() && !detected.await(Seconds(10))) {
    return Error("Failed to wait " + stringify(Seconds(10)) +
                 " for master to detect the leader");
  }

  return master;
}

Try<process::PID<master::Master> > MesosTest::StartMaster(
    master::allocator::AllocatorProcess* allocator,
    const Option<master::Flags>& flags,
    bool wait)
{
  Future<Nothing> detected = FUTURE_DISPATCH(_, &master::Master::detected);

  Try<process::PID<master::Master> > master = cluster.masters.start(
      allocator, flags.isNone() ? CreateMasterFlags() : flags.get());

  // Wait until the leader is detected because otherwise this master
  // may reject authentication requests because it doesn't know it's
  // the leader yet [MESOS-881].
  if (wait && master.isSome() && !detected.await(Seconds(10))) {
    return Error("Failed to wait " + stringify(Seconds(10)) +
                 " for master to detect the leader");
  }

  return master;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer();

  Try<process::PID<slave::Slave> > pid = StartSlave(containerizer, flags);

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    MockExecutor* executor,
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer(executor);

  Try<process::PID<slave::Slave> > pid = StartSlave(containerizer, flags);

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    slave::Containerizer* containerizer,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      containerizer, flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    slave::Containerizer* containerizer,
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      containerizer,
      detector,
      flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<PID<slave::Slave> > MesosTest::StartSlave(
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      detector, flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<PID<slave::Slave> > MesosTest::StartSlave(
    MockExecutor* executor,
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer(executor);

  Try<process::PID<slave::Slave> > pid = cluster.slaves.start(
      containerizer,
      detector,
      flags.isNone() ? CreateSlaveFlags() : flags.get());

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

void MesosTest::Stop(const process::PID<master::Master>& pid)
{
  cluster.masters.stop(pid);
}

void MesosTest::Stop(const process::PID<slave::Slave>& pid, bool shutdown)
{
  cluster.slaves.stop(pid, shutdown);
  if (containerizers.count(pid) > 0) {
    slave::Containerizer* containerizer = containerizers[pid];
    containerizers.erase(pid);
    delete containerizer;
  }
}

void MesosTest::Shutdown()
{
  ShutdownMasters();
  ShutdownSlaves();
}

void MesosTest::ShutdownMasters()
{
  cluster.masters.shutdown();
}

void MesosTest::ShutdownSlaves()
{
  cluster.slaves.shutdown();

  foreachvalue (slave::Containerizer* containerizer, containerizers) {
    delete containerizer;
  }
  containerizers.clear();
}

void MesosTest::TearDown()
{
  // TODO(benh): Fail the test if shutdown hasn't been called?
  Shutdown();
}

slave::Flags ContainerizerTest<slave::MesosContainerizer>::CreateSlaveFlags()
{
  slave::Flags flags = MesosTest::CreateSlaveFlags();

#ifdef __linux__
  flags.isolation = "cgroups/cpu,cgroups/mem";
  flags.cgroups_hierarchy = baseHierarchy;
  flags.cgroups_root = TEST_CGROUPS_ROOT + "_" + UUID::random().toString();
#else
  flags.isolation = "posix/cpu,posix/mem";
#endif

  return flags;
}

#ifdef __linux__
void ContainerizerTest<slave::MesosContainerizer>::SetUpTestCase()
{
  // Clean up any testing hierarchies.
  Try<std::set<std::string> > hierarchies = cgroups::hierarchies();
  ASSERT_SOME(hierarchies);
  foreach (const std::string& hierarchy, hierarchies.get()) {
    if (strings::startsWith(hierarchy, TEST_CGROUPS_HIERARCHY)) {
      AWAIT_READY(cgroups::cleanup(hierarchy));
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::TearDownTestCase()
{
  // Clean up any testing hierarchies.
  Try<std::set<std::string> > hierarchies = cgroups::hierarchies();
  ASSERT_SOME(hierarchies);
  foreach (const std::string& hierarchy, hierarchies.get()) {
    if (strings::startsWith(hierarchy, TEST_CGROUPS_HIERARCHY)) {
      AWAIT_READY(cgroups::cleanup(hierarchy));
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::SetUp()
{
  MesosTest::SetUp();

  subsystems.insert("cpu");
  subsystems.insert("cpuacct");
  subsystems.insert("memory");
  subsystems.insert("freezer");

  foreach (const std::string& subsystem, subsystems) {
    // Establish the base hierarchy if this is the first subsystem checked.
    if (baseHierarchy.empty()) {
      Result<std::string> hierarchy = cgroups::hierarchy(subsystem);
      ASSERT_FALSE(hierarchy.isError());

      if (hierarchy.isNone()) {
        baseHierarchy = TEST_CGROUPS_HIERARCHY;
      } else {
        // Strip the subsystem to get the base hierarchy.
        baseHierarchy = strings::remove(
            hierarchy.get(),
            subsystem,
            strings::SUFFIX);
      }
    }

    // Mount the subsystem if necessary.
    std::string hierarchy = path::join(baseHierarchy, subsystem);
    Try<bool> mounted = cgroups::mounted(hierarchy, subsystem);
    ASSERT_SOME(mounted);
    if (!mounted.get()) {
      ASSERT_SOME(cgroups::mount(hierarchy, subsystem))
        << "-------------------------------------------------------------\n"
        << "We cannot run any cgroups tests that require\n"
        << "a hierarchy with subsystem '" << subsystem << "'\n"
        << "because we failed to find an existing hierarchy\n"
        << "or create a new one (tried '" << hierarchy << "').\n"
        << "You can either remove all existing\n"
        << "hierarchies, or disable this test case\n"
        << "(i.e., --gtest_filter=-"
        << ::testing::UnitTest::GetInstance()
            ->current_test_info()
            ->test_case_name() << ".*).\n"
        << "-------------------------------------------------------------";
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::TearDown()
{
  MesosTest::TearDown();

  foreach (const std::string& subsystem, subsystems) {
    std::string hierarchy = path::join(baseHierarchy, subsystem);

    Try<std::vector<std::string> > cgroups = cgroups::get(hierarchy);
    CHECK_SOME(cgroups);

    foreach (const std::string& cgroup, cgroups.get()) {
      // Remove any cgroups that start with TEST_CGROUPS_ROOT
      if (strings::startsWith(cgroup, TEST_CGROUPS_ROOT)) {
        AWAIT_READY(cgroups::destroy(hierarchy, cgroup));
      }
    }
  }
}
#endif // __linux__

} // namespace tests {
} // namespace internal {
} // namespace mesos {
<commit_msg>Fixed tests to not use cgroup isolators unless available and running as root.<commit_after>#include <stout/check.hpp>
#include <stout/foreach.hpp>
#include <stout/os.hpp>
#include <stout/path.hpp>
#include <stout/result.hpp>
#include <stout/uuid.hpp>

#ifdef __linux__
#include "linux/cgroups.hpp"
#endif

#include "slave/containerizer/containerizer.hpp"
#include "slave/containerizer/mesos_containerizer.hpp"

#include "tests/containerizer.hpp"
#include "tests/environment.hpp"
#include "tests/flags.hpp"
#include "tests/mesos.hpp"

using namespace process;

using testing::_;

namespace mesos {
namespace internal {
namespace tests {

MesosTest::MesosTest(const Option<zookeeper::URL>& url) : cluster(url) {}

master::Flags MesosTest::CreateMasterFlags()
{
  master::Flags flags;

  // Enable authentication.
  flags.authenticate = true;

  // Create a temporary work directory (removed by Environment).
  const Try<std::string>& directory = environment->mkdtemp();
  CHECK_SOME(directory) << "Failed to create temporary directory";

  flags.work_dir = directory.get();

  // Create a default credentials file.
  const std::string& path = path::join(directory.get(), "credentials");

  const std::string& credentials =
    DEFAULT_CREDENTIAL.principal() + " " + DEFAULT_CREDENTIAL.secret();

  CHECK_SOME(os::write(path, credentials))
    << "Failed to write credentials to '" << path << "'";

  flags.credentials = "file://" + path;

  return flags;
}

slave::Flags MesosTest::CreateSlaveFlags()
{
  slave::Flags flags;

  // Create a temporary work directory (removed by Environment).
  Try<std::string> directory = environment->mkdtemp();
  CHECK_SOME(directory) << "Failed to create temporary directory";

  flags.work_dir = directory.get();

  flags.launcher_dir = path::join(tests::flags.build_dir, "src");

  // TODO(vinod): Consider making this true and fixing the tests.
  flags.checkpoint = false;

  flags.resources = Option<std::string>(
      "cpus:2;mem:1024;disk:1024;ports:[31000-32000]");

  return flags;
}

Try<process::PID<master::Master> > MesosTest::StartMaster(
    const Option<master::Flags>& flags,
    bool wait)
{
  Future<Nothing> detected = FUTURE_DISPATCH(_, &master::Master::detected);

  Try<process::PID<master::Master> > master = cluster.masters.start(
      flags.isNone() ? CreateMasterFlags() : flags.get());

  // Wait until the leader is detected because otherwise this master
  // may reject authentication requests because it doesn't know it's
  // the leader yet [MESOS-881].
  if (wait && master.isSome() && !detected.await(Seconds(10))) {
    return Error("Failed to wait " + stringify(Seconds(10)) +
                 " for master to detect the leader");
  }

  return master;
}

Try<process::PID<master::Master> > MesosTest::StartMaster(
    master::allocator::AllocatorProcess* allocator,
    const Option<master::Flags>& flags,
    bool wait)
{
  Future<Nothing> detected = FUTURE_DISPATCH(_, &master::Master::detected);

  Try<process::PID<master::Master> > master = cluster.masters.start(
      allocator, flags.isNone() ? CreateMasterFlags() : flags.get());

  // Wait until the leader is detected because otherwise this master
  // may reject authentication requests because it doesn't know it's
  // the leader yet [MESOS-881].
  if (wait && master.isSome() && !detected.await(Seconds(10))) {
    return Error("Failed to wait " + stringify(Seconds(10)) +
                 " for master to detect the leader");
  }

  return master;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer();

  Try<process::PID<slave::Slave> > pid = StartSlave(containerizer, flags);

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    MockExecutor* executor,
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer(executor);

  Try<process::PID<slave::Slave> > pid = StartSlave(containerizer, flags);

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    slave::Containerizer* containerizer,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      containerizer, flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<process::PID<slave::Slave> > MesosTest::StartSlave(
    slave::Containerizer* containerizer,
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      containerizer,
      detector,
      flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<PID<slave::Slave> > MesosTest::StartSlave(
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  return cluster.slaves.start(
      detector, flags.isNone() ? CreateSlaveFlags() : flags.get());
}

Try<PID<slave::Slave> > MesosTest::StartSlave(
    MockExecutor* executor,
    Owned<MasterDetector> detector,
    const Option<slave::Flags>& flags)
{
  slave::Containerizer* containerizer = new TestContainerizer(executor);

  Try<process::PID<slave::Slave> > pid = cluster.slaves.start(
      containerizer,
      detector,
      flags.isNone() ? CreateSlaveFlags() : flags.get());

  if (pid.isError()) {
    delete containerizer;
    return pid;
  }

  containerizers[pid.get()] = containerizer;

  return pid;
}

void MesosTest::Stop(const process::PID<master::Master>& pid)
{
  cluster.masters.stop(pid);
}

void MesosTest::Stop(const process::PID<slave::Slave>& pid, bool shutdown)
{
  cluster.slaves.stop(pid, shutdown);
  if (containerizers.count(pid) > 0) {
    slave::Containerizer* containerizer = containerizers[pid];
    containerizers.erase(pid);
    delete containerizer;
  }
}

void MesosTest::Shutdown()
{
  ShutdownMasters();
  ShutdownSlaves();
}

void MesosTest::ShutdownMasters()
{
  cluster.masters.shutdown();
}

void MesosTest::ShutdownSlaves()
{
  cluster.slaves.shutdown();

  foreachvalue (slave::Containerizer* containerizer, containerizers) {
    delete containerizer;
  }
  containerizers.clear();
}

void MesosTest::TearDown()
{
  // TODO(benh): Fail the test if shutdown hasn't been called?
  Shutdown();
}

slave::Flags ContainerizerTest<slave::MesosContainerizer>::CreateSlaveFlags()
{
  slave::Flags flags = MesosTest::CreateSlaveFlags();

#ifdef __linux__
  // Use cgroup isolators if they're available and we're root.
  // TODO(idownes): Refactor the cgroups/non-cgroups code.
  if (os::exists("/proc/cgroups") && os::user() == "root") {
    flags.isolation = "cgroups/cpu,cgroups/mem";
    flags.cgroups_hierarchy = baseHierarchy;
    flags.cgroups_root = TEST_CGROUPS_ROOT + "_" + UUID::random().toString();
  } else {
    flags.isolation = "posix/cpu,posix/mem";
  }
#else
  flags.isolation = "posix/cpu,posix/mem";
#endif

  return flags;
}

#ifdef __linux__
void ContainerizerTest<slave::MesosContainerizer>::SetUpTestCase()
{
  if (os::exists("/proc/cgroups") && os::user() == "root") {
    // Clean up any testing hierarchies.
    Try<std::set<std::string> > hierarchies = cgroups::hierarchies();
    ASSERT_SOME(hierarchies);
    foreach (const std::string& hierarchy, hierarchies.get()) {
      if (strings::startsWith(hierarchy, TEST_CGROUPS_HIERARCHY)) {
        AWAIT_READY(cgroups::cleanup(hierarchy));
      }
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::TearDownTestCase()
{
  if (os::exists("/proc/cgroups") && os::user() == "root") {
    // Clean up any testing hierarchies.
    Try<std::set<std::string> > hierarchies = cgroups::hierarchies();
    ASSERT_SOME(hierarchies);
    foreach (const std::string& hierarchy, hierarchies.get()) {
      if (strings::startsWith(hierarchy, TEST_CGROUPS_HIERARCHY)) {
        AWAIT_READY(cgroups::cleanup(hierarchy));
      }
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::SetUp()
{
  MesosTest::SetUp();

  subsystems.insert("cpu");
  subsystems.insert("cpuacct");
  subsystems.insert("memory");
  subsystems.insert("freezer");

  if (os::exists("/proc/cgroups") && os::user() == "root") {
    foreach (const std::string& subsystem, subsystems) {
      // Establish the base hierarchy if this is the first subsystem checked.
      if (baseHierarchy.empty()) {
        Result<std::string> hierarchy = cgroups::hierarchy(subsystem);
        ASSERT_FALSE(hierarchy.isError());

        if (hierarchy.isNone()) {
          baseHierarchy = TEST_CGROUPS_HIERARCHY;
        } else {
          // Strip the subsystem to get the base hierarchy.
          baseHierarchy = strings::remove(
              hierarchy.get(),
              subsystem,
              strings::SUFFIX);
        }
      }

      // Mount the subsystem if necessary.
      std::string hierarchy = path::join(baseHierarchy, subsystem);
      Try<bool> mounted = cgroups::mounted(hierarchy, subsystem);
      ASSERT_SOME(mounted);
      if (!mounted.get()) {
        ASSERT_SOME(cgroups::mount(hierarchy, subsystem))
          << "-------------------------------------------------------------\n"
          << "We cannot run any cgroups tests that require\n"
          << "a hierarchy with subsystem '" << subsystem << "'\n"
          << "because we failed to find an existing hierarchy\n"
          << "or create a new one (tried '" << hierarchy << "').\n"
          << "You can either remove all existing\n"
          << "hierarchies, or disable this test case\n"
          << "(i.e., --gtest_filter=-"
          << ::testing::UnitTest::GetInstance()
              ->current_test_info()
              ->test_case_name() << ".*).\n"
          << "-------------------------------------------------------------";
      }
    }
  }
}

void ContainerizerTest<slave::MesosContainerizer>::TearDown()
{
  MesosTest::TearDown();

  if (os::exists("/proc/cgroups") && os::user() == "root") {
    foreach (const std::string& subsystem, subsystems) {
      std::string hierarchy = path::join(baseHierarchy, subsystem);

      Try<std::vector<std::string> > cgroups = cgroups::get(hierarchy);
      CHECK_SOME(cgroups);

      foreach (const std::string& cgroup, cgroups.get()) {
        // Remove any cgroups that start with TEST_CGROUPS_ROOT
        if (strings::startsWith(cgroup, TEST_CGROUPS_ROOT)) {
          AWAIT_READY(cgroups::destroy(hierarchy, cgroup));
        }
      }
    }
  }
}
#endif // __linux__

} // namespace tests {
} // namespace internal {
} // namespace mesos {
<|endoftext|>