modeco80
/
singrdk
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
singrdk/base/Windows/JobControl/jobcontrol.c

699 lines
21 KiB
C
Raw Blame History

// ----------------------------------------------------------------------------
//
// Program: JobControl
//
// Purpose: Run multiple processes within a job object and/or manually kill
// existing job objects. See usage for more details.
//
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// ----------------------------------------------------------------------------
#define _WIN32_WINNT 0x0500
#define YESWINDOWSTATION
#define YESSECURITY
#include <winlean.h>
#include <assert.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
typedef struct _CONTROLLED_PROCESS {
HANDLE Thread;
HANDLE Process;
DWORD ProcessId;
BOOL ErrorExitFatal;
} CONTROLLED_PROCESS, *PCONTROLLED_PROCESS;
typedef struct _CONTROLLED_PROCESS_SET {
DWORD ProcessCount;
DWORD MaxProcessCount;
PCONTROLLED_PROCESS Processes;
} CONTROLLED_PROCESS_SET, *PCONTROLLED_PROCESS_SET;
typedef struct _MONITOR_THREAD_ARGS {
HANDLE hJob;
PCONTROLLED_PROCESS_SET pSet;
DWORD dwTimeoutSeconds;
HANDLE hMonitoringStarted;
} MONITOR_THREAD_ARGS, *PMONITOR_THREAD_ARGS;
// ----------------------------------------------------------------------------
static const DWORD dwCompletionKey = 2;
// Global job handle needed for graceful shutdown by control
// event handler.
static HANDLE ghJob = NULL;
static BOOL gbVerbose = FALSE;
#define tprintf(x, ...) \
do { \
if (gbVerbose) printf(x, __VA_ARGS__); \
} while (FALSE)
// ----------------------------------------------------------------------------
static BOOL WINAPI CtrlHandler(DWORD dwCtrlType)
{
tprintf("Caught control event %I32u", dwCtrlType);
switch (dwCtrlType) {
case CTRL_C_EVENT:
case CTRL_BREAK_EVENT:
GenerateConsoleCtrlEvent(dwCtrlType, 0);
default:
break;
}
if (ghJob != NULL) {
TerminateJobObject(ghJob, (UINT)-1);
}
ExitProcess((UINT)-2);
}
static void Usage()
{
printf(
"Usage:\n" \
"jobcontrol create|kill <jobname> [options]\n" \
"\n" \
"where command is one of create, kill, watch.\n" \
"\n" \
"CREATE COMMAND\n" \
"--------------\n" \
"\n" \
"The \"create\" command creates a job object and runs any number\n" \
"of commands with it. The command runs for as long as the\n" \
"specified commands are executing.\n" \
"\n" \
"The options available are:\n" \
" /LimitJobTime <seconds> - Set limit on elapsed job time.\n" \
" /LimitJobUserTime <seconds> - Set limit on elapsed job user mode time.\n" \
" /LimitProcessUserTime <seconds> - Set limit on per-process user mode time\n" \
" /C <command> - Run command in job.\n" \
" /T <command> - Run command in job and terminate if\n" \
" the command fails (exit code != 0).\n" \
" /V - Run verbosely.\n" \
"\n" \
"Example:\n" \
" jobcontrol create MyJob /LimitJobTime 10 /C \"cmd.exe\"\n" \
"\n" \
"KILL COMMAND\n" \
"------------\n" \
"\n" \
"The \"kill\" command terminates the processes within a job\n" \
"object. There are no relevant options.\n" \
"\n" \
"Example:\n" \
" jobcontrol kill MyJob\n" \
"\n"
);
}
static void
ErrorReport(char *function)
{
VOID *message;
DWORD error = GetLastError();
FormatMessage(
FORMAT_MESSAGE_ALLOCATE_BUFFER |
FORMAT_MESSAGE_FROM_SYSTEM,
NULL,
error,
MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
(LPTSTR) &message,
0, NULL );
printf("jobcontrol: %s failed with error %d: %s",
function, error, message);
LocalFree(message);
}
static BOOL
ParseInt64(LPCSTR lpszValue, PINT64 pValue)
{
PCHAR pszStop;
errno = 0;
*pValue = _strtoi64(lpszValue, &pszStop, 10);
// Return no overflow and all characters parsed.
return errno == 0 && *pszStop == '\0';
}
static BOOL
ParseUInt32(LPCSTR lpszValue, PUINT32 pValue)
{
PCHAR pszStop;
errno = 0;
*pValue = strtoul(lpszValue, &pszStop, 10);
// Return no overflow and all characters parsed.
return errno == 0 && *pszStop == '\0';
}
// ----------------------------------------------------------------------------
// Interactivity check
static BOOL IsInteractive()
{
static int isInteractive = -1;
if (isInteractive < 0) {
HANDLE hwinsta;
if (NULL != (hwinsta = GetProcessWindowStation())) {
USEROBJECTFLAGS uof;
DWORD dwLengthNeeded;
if (GetUserObjectInformation(hwinsta, UOI_FLAGS, &uof,
sizeof(uof), &dwLengthNeeded) == 0) {
ErrorReport("GetUserObjectInformation");
ExitProcess((DWORD)-1);
}
isInteractive = (uof.dwFlags & WSF_VISIBLE) ? 1 : 0;
tprintf("Process is interactive: %c\n",
isInteractive ? 't' : 'f');
}
else {
ErrorReport("GetProcessWindowStation");
ExitProcess((DWORD)-1);
}
}
return isInteractive ? TRUE : FALSE;
}
// ----------------------------------------------------------------------------
// Process Set helpers
static PCONTROLLED_PROCESS_SET
AllocateProcessSet(DWORD dwMaxProcesses)
{
size_t cbBytes = (sizeof(CONTROLLED_PROCESS_SET) +
dwMaxProcesses * sizeof(CONTROLLED_PROCESS));
PCONTROLLED_PROCESS_SET pSet = (PCONTROLLED_PROCESS_SET)malloc(cbBytes);
if (NULL != pSet) {
pSet->ProcessCount = 0;
pSet->MaxProcessCount = dwMaxProcesses;
pSet->Processes = (PCONTROLLED_PROCESS)(pSet + 1);
}
return pSet;
}
static void
FreeProcessSet(PCONTROLLED_PROCESS_SET pSet)
{
free(pSet);
pSet = NULL;
}
static PCONTROLLED_PROCESS
FindProcessInSet(PCONTROLLED_PROCESS_SET pSet, DWORD dwProcessId)
{
DWORD i;
for (i = 0; i < pSet->ProcessCount; i++) {
if (pSet->Processes[i].ProcessId == dwProcessId) {
return &pSet->Processes[i];
}
}
return NULL;
}
static void
AddProcessToSet(PCONTROLLED_PROCESS_SET pSet,
HANDLE hThread,
HANDLE hProcess,
DWORD dwProcessId,
BOOL bErrorExitFatal)
{
PCONTROLLED_PROCESS pProcess;
assert(FindProcessInSet(pSet, dwProcessId) == NULL);
assert(pSet->ProcessCount < pSet->MaxProcessCount);
pProcess = &pSet->Processes[pSet->ProcessCount];
pProcess->Thread = hThread;
pProcess->Process = hProcess;
pProcess->ProcessId = dwProcessId;
pProcess->ErrorExitFatal = bErrorExitFatal;
pSet->ProcessCount++;
}
static BOOL
RemoveProcessFromSet(PCONTROLLED_PROCESS_SET pSet, DWORD dwProcessId)
{
DWORD i;
for (i = 0; i < pSet->ProcessCount; i++) {
if (pSet->Processes[i].ProcessId == dwProcessId) {
tprintf("Removing process %u from process set\n", dwProcessId);
if (pSet->ProcessCount > 1) {
// Swap with last element so it's there for inspection
CONTROLLED_PROCESS tmp = pSet->Processes[i];
pSet->Processes[i] = pSet->Processes[pSet->ProcessCount - 1];
pSet->Processes[pSet->ProcessCount - 1] = tmp;
}
pSet->ProcessCount--;
return TRUE;
}
}
return FALSE;
}
// ----------------------------------------------------------------------------
static BOOL
LimitJobUserTime(HANDLE hJob, DWORD dwLimitFlags, LPCSTR lpszValue)
{
JOBOBJECT_BASIC_LIMIT_INFORMATION limits;
__int64 timeout;
__asm int 3;
ZeroMemory(&limits, sizeof(limits));
if (ParseInt64(lpszValue, &timeout) == FALSE) {
fprintf(stdout, "Bad user time limit value.");
return FALSE;
}
tprintf("Limiting job time to %I64d seconds\n", timeout);
timeout *= 1000 * 1000 * 10; // seconds to W32 kernel ticks
tprintf("Limiting job time to %I64d ticks\n", timeout);
if (dwLimitFlags == JOB_OBJECT_LIMIT_JOB_TIME) {
CopyMemory(&limits.PerJobUserTimeLimit, &timeout, sizeof(timeout));
}
else if (dwLimitFlags == JOB_OBJECT_LIMIT_PROCESS_TIME) {
CopyMemory(&limits.PerProcessUserTimeLimit, &timeout, sizeof(timeout));
}
else {
fprintf(stdout, "Bad time limit target.\n");
return FALSE;
}
limits.LimitFlags = dwLimitFlags;
if (SetInformationJobObject(hJob, JobObjectBasicLimitInformation,
&limits, sizeof(limits)) == 0) {
ErrorReport("SetInformationJobObject");
return FALSE;
}
ErrorReport("SetInformationJobObject (no error)");
return TRUE;
}
static BOOL
AddProcessToJob(HANDLE hJob,
LPCSTR lpCommandLine,
PCONTROLLED_PROCESS_SET pSet,
DWORD bErrorExitFatal)
{
PROCESS_INFORMATION pi;
STARTUPINFO si;
DWORD dwCreationFlags;
ZeroMemory(&si, sizeof(si));
si.cb = sizeof(si);
dwCreationFlags = CREATE_SUSPENDED /* | CREATE_BREAKAWAY_FROM_JOB */;
if (!IsInteractive()) {
// If non-interactive we are presumably being run by a
// process and do not want to interact with the desktop.
// If we attempt to create a window, CreateProcess will
// die with permission denied.
dwCreationFlags |= CREATE_NO_WINDOW;
}
if (!CreateProcess(NULL, (LPSTR)lpCommandLine, NULL, NULL, FALSE,
dwCreationFlags, NULL, NULL, &si, &pi)) {
fprintf(stdout, "Failed to create process: \"%s\"\n", lpCommandLine);
ErrorReport("CreateProcess");
return FALSE;
}
else if (!AssignProcessToJobObject(hJob, pi.hProcess)) {
ErrorReport("AssignProcessToJobObject");
TerminateProcess(pi.hProcess, (UINT)-1);
return FALSE;
}
AddProcessToSet(pSet, pi.hThread, pi.hProcess, pi.dwProcessId,
bErrorExitFatal);
printf("jobcontrol: Adding process (%u): %s\n",
pi.dwProcessId, lpCommandLine);
return TRUE;
}
static BOOL
CreateAndAssociateCompletionPort(HANDLE hJob, PHANDLE phIoPort)
{
JOBOBJECT_ASSOCIATE_COMPLETION_PORT jacp;
HANDLE hIoPort;
DWORD dwKey;
*phIoPort = INVALID_HANDLE_VALUE;
if (NULL == (hIoPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE,
NULL, dwCompletionKey, 1))) {
ErrorReport("CreateIoCompletionPort");
return FALSE;
}
dwKey = dwCompletionKey;
jacp.CompletionKey = (PVOID)dwKey;
jacp.CompletionPort = hIoPort;
if (!SetInformationJobObject(hJob,
JobObjectAssociateCompletionPortInformation,
&jacp, sizeof(jacp))) {
ErrorReport("SetInformationJobObject");
CloseHandle(hIoPort);
return FALSE;
}
*phIoPort = hIoPort;
return TRUE;
}
static DWORD WINAPI MonitorJobThreadProc(LPVOID lpParams)
{
PMONITOR_THREAD_ARGS pJobArgs = (PMONITOR_THREAD_ARGS)lpParams;
HANDLE hIoPort = NULL;
BOOL bStop = FALSE;
BOOL bFatalError = FALSE;
time_t nowTime = 0i64;
time_t endTime = _I64_MAX;
if (!CreateAndAssociateCompletionPort(pJobArgs->hJob, &hIoPort)) {
ExitThread((DWORD)-1);
}
if (pJobArgs->dwTimeoutSeconds > 0) {
time(&nowTime);
endTime = nowTime + pJobArgs->dwTimeoutSeconds;
}
tprintf("Watching job (supplied = %u, %I64d, %I64d, timeout = %I64d)\n",
pJobArgs->dwTimeoutSeconds, endTime, nowTime, endTime - nowTime);
SetEvent(pJobArgs->hMonitoringStarted);
while (time(&nowTime) < endTime && !bStop) {
DWORD dwBytes = 0;
ULONG_PTR key = 0;
LPOVERLAPPED pOverlapped = NULL;
DWORD waitMillis = INFINITE;
if (endTime - nowTime < INFINITE * 1000i64) {
waitMillis = (DWORD)((endTime - nowTime) * 1000);
}
if (!GetQueuedCompletionStatus(hIoPort, &dwBytes, &key,
&pOverlapped, waitMillis)) {
continue;
}
assert (key == dwCompletionKey);
switch (dwBytes) {
case JOB_OBJECT_MSG_END_OF_JOB_TIME:
case JOB_OBJECT_MSG_END_OF_PROCESS_TIME:
case JOB_OBJECT_MSG_ACTIVE_PROCESS_LIMIT:
case JOB_OBJECT_MSG_PROCESS_MEMORY_LIMIT:
case JOB_OBJECT_MSG_JOB_MEMORY_LIMIT:
break;
case JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO:
printf("jobcontrol: Requested process all stopped.");
bStop = TRUE;
break;
case JOB_OBJECT_MSG_NEW_PROCESS:
{
DWORD dwProcessId = (DWORD)pOverlapped;
tprintf("Process start (%d)\n", dwProcessId);
}
break;
case JOB_OBJECT_MSG_EXIT_PROCESS:
case JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS:
{
DWORD dwExitCode = 0;
DWORD dwProcessId = (DWORD)pOverlapped;
PCONTROLLED_PROCESS pProcess =
FindProcessInSet(pJobArgs->pSet, dwProcessId);
if (NULL != pProcess) {
if (!GetExitCodeProcess(pProcess->Process, &dwExitCode)) {
ErrorReport("GetExitCodeProcess");
fprintf(stdout, "jobcontrol: Unable to get exit code of process in job\n");
bStop = TRUE;
bFatalError = TRUE;
}
else if (dwExitCode != 0 && pProcess->ErrorExitFatal) {
bStop = TRUE;
bFatalError = TRUE;
}
printf("jobcontrol: Process %u exited (%u)\n",
dwProcessId, dwExitCode);
RemoveProcessFromSet(pJobArgs->pSet, dwProcessId);
if (pJobArgs->pSet->ProcessCount == 0) {
bStop = TRUE;
}
}
}
break;
default:
// Should never happen
printf("jobcontrol: Unknown message type %08x", dwBytes);
break;
}
}
if (nowTime > endTime) {
printf("jobcontrol: Timed out\n");
}
tprintf("- Watching job (supplied = %u, %I64d, %I64d, timeout = %I64d)\n",
pJobArgs->dwTimeoutSeconds, endTime, nowTime, endTime - nowTime);
tprintf("Watch done (fatal error = %d, stop = %d)\n", bFatalError, bStop);
CloseHandle(hIoPort);
ExitThread((bFatalError || !bStop) ? (DWORD)-1 : 0);
}
static BOOL
StartJobThreads(PCONTROLLED_PROCESS_SET pSet)
{
DWORD i;
for (i = 0; i < pSet->ProcessCount; i++) {
tprintf("Starting process %d\n", pSet->Processes[i].ProcessId);
ResumeThread(pSet->Processes[i].Thread);
pSet->Processes[i].Thread = NULL;
}
return TRUE;
}
static void
CleanupJobHandles(PCONTROLLED_PROCESS_SET pSet)
{
DWORD i;
PCONTROLLED_PROCESS pProcess = pSet->Processes;
for (i = 0; i < pSet->MaxProcessCount && pProcess->ProcessId != 0;
i++, pProcess++) {
CloseHandle(pProcess->Thread);
CloseHandle(pProcess->Process);
}
}
static DWORD
StartJobProcessesAndMonitor(HANDLE hJob,
PCONTROLLED_PROCESS_SET pSet,
DWORD dwTimeoutSeconds)
{
HANDLE hMonitorThread;
DWORD dwMonitorThreadExitCode;
MONITOR_THREAD_ARGS monitorArgs;
//
// Fire up monitoring thread so process can track all activity
//
monitorArgs.hJob = hJob;
monitorArgs.pSet = pSet;
monitorArgs.dwTimeoutSeconds = dwTimeoutSeconds;
monitorArgs.hMonitoringStarted = CreateEvent(NULL, TRUE, FALSE, NULL);
hMonitorThread = CreateThread(NULL, 0, MonitorJobThreadProc,
&monitorArgs, 0, NULL);
if (NULL == hMonitorThread) {
CloseHandle(monitorArgs.hMonitoringStarted);
ErrorReport("CreateThread");
return (DWORD)-1;
}
//
// Wait for monitoring thread to reach start point
//
WaitForSingleObject(monitorArgs.hMonitoringStarted, INFINITE);
//
// Resume process threads associated with job
//
StartJobThreads(pSet);
//
// Wait for monitoring thread to kick the bucket
//
WaitForSingleObject(hMonitorThread, INFINITE);
GetExitCodeThread(hMonitorThread, &dwMonitorThreadExitCode);
CloseHandle(hMonitorThread);
CloseHandle(monitorArgs.hMonitoringStarted);
CleanupJobHandles(pSet);
return dwMonitorThreadExitCode;
}
// ----------------------------------------------------------------------------
// User Commands
static int
CreateJobAndRun(int argc, LPCSTR argv[])
{
BOOL bSuccess, bInJob;
DWORD dwMaxJobTime = 0;
DWORD dwResult;
PCONTROLLED_PROCESS_SET pSet;
if (argc == 0) {
fprintf(stdout, "Missing <jobname> and options\n");
return -1;
}
if (!IsProcessInJob(GetCurrentProcess(), NULL, &bInJob)) {
ErrorReport("IsProcessInJob");
}
else if (bInJob) {
fprintf(stdout, "Warning: jobcontrol launched inside a job\n");
}
tprintf("Creating job %s\n", argv[0]);
if (NULL == (ghJob = CreateJobObject(NULL, argv[0]))) {
ErrorReport("CreateJobObject");
return -1;
}
argv++;
argc--;
pSet = AllocateProcessSet(argc); // overallocation
bSuccess = TRUE;
while (argc != 0 && bSuccess) {
if (argc < 2) {
fprintf(stdout, "Unexpected end of arguments\n");
bSuccess = FALSE;
}
else if (!_strcmpi("/LimitJobUserTime", argv[0])) {
bSuccess = LimitJobUserTime(ghJob, JOB_OBJECT_LIMIT_JOB_TIME,
argv[1]);
}
else if (!_strcmpi("/LimitProcessUserTime", argv[0])) {
bSuccess = LimitJobUserTime(ghJob,JOB_OBJECT_LIMIT_PROCESS_TIME,
argv[1]);
}
else if (!_strcmpi("/LimitJobTime", argv[0])) {
bSuccess = ParseUInt32(argv[1], (PUINT32)&dwMaxJobTime);
}
else if (!_strcmpi("/C", argv[0])) {
bSuccess = AddProcessToJob(ghJob, argv[1], pSet, FALSE);
}
else if (!_strcmpi("/T", argv[0])) {
bSuccess = AddProcessToJob(ghJob, argv[1], pSet, TRUE);
}
else if (!_strcmpi("/V", argv[0])) {
gbVerbose = TRUE;
argv += 1;
argc -= 1;
continue;
}
else {
fprintf(stdout, "Unknown option \"%s\".\n", argv[0]);
bSuccess = FALSE;
}
argv += 2;
argc -= 2;
}
dwResult = (DWORD)-1;
if (bSuccess && pSet->ProcessCount != 0) {
dwResult = StartJobProcessesAndMonitor(ghJob, pSet, dwMaxJobTime);
}
FreeProcessSet(pSet);
TerminateJobObject(ghJob, (UINT)-1);
CloseHandle(ghJob);
return (int)dwResult;
}
static int
KillJob(int argc, LPCSTR argv[])
{
HANDLE hJob;
if (argc == 0) {
fprintf(stdout, "Missing <jobname> argument\n");
return -1;
}
else if (argc > 1) {
fprintf(stdout, "Unknown arguments passed to kill request\n");
Usage();
return -1;
}
else if (NULL == (hJob = CreateJobObject(NULL, argv[0]))) {
ErrorReport("CreateJobObject");
return -1;
}
else if (0 == TerminateJobObject(hJob, (UINT)-1)) {
ErrorReport("TerminateJobObject");
return -1;
}
tprintf("Terminated %s", argv[0]);
return 0;
}
// ----------------------------------------------------------------------------
// Main
int main(int argc, const char* argv[])
{
if (argc >= 3) {
if (_strcmpi(argv[1], "create") == 0) {
SetConsoleCtrlHandler(CtrlHandler, TRUE);
return CreateJobAndRun(argc - 2, argv + 2);
}
else if (_strcmpi(argv[1], "kill") == 0) {
SetConsoleCtrlHandler(CtrlHandler, TRUE);
return KillJob(argc - 2, argv + 2);
}
else {
Usage();
}
}
else {
Usage();
return -1;
}
}
Reference in New Issue View Git Blame Copy Permalink