Skip to content

Commit

Permalink
Refactor Process.open_files() to minimize threaded section
Browse files Browse the repository at this point in the history 
Signed-off-by: Jesse Schwartzentruber <truber@mozilla.com>
  • Loading branch information
@jschwartzentruber
jschwartzentruber committed Jan 10, 2023
1 parent 2da9950 commit b060339
Showing 1 changed file with 108 additions and 124 deletions.
232 changes: 108 additions & 124 deletions psutil/arch/windows/process_handles.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,15 @@


#define THREAD_TIMEOUT 100 // ms
// Global object shared between the 2 threads.
PUNICODE_STRING globalFileName = NULL;
#define STATUS_UNSUCCESSFUL ((DWORD)0xC0000001L)

typedef struct
{
NTSTATUS status;
HANDLE hFile;
PUNICODE_STRING FileName;
ULONG Length;
} psutil_query_thread;


static int
Expand Down Expand Up @@ -80,129 +87,105 @@ psutil_enum_handles(PSYSTEM_HANDLE_INFORMATION_EX *handles) {
return 0;
}


static int
static DWORD
psutil_get_filename(LPVOID lpvParam) {
HANDLE hFile = *((HANDLE*)lpvParam);
NTSTATUS status;
ULONG bufferSize;
ULONG attempts = 8;

bufferSize = 0x200;
globalFileName = MALLOC_ZERO(bufferSize);
if (globalFileName == NULL) {
PyErr_NoMemory();
goto error;
}
psutil_query_thread* context = lpvParam;

context->status = NtQueryObject(
context->hFile,
ObjectNameInformation,
context->FileName,
context->Length,
&context->Length);
return 0;
}


// Note: also this is supposed to hang, hence why we do it in here.
if (GetFileType(hFile) != FILE_TYPE_DISK) {
SetLastError(0);
globalFileName->Length = 0;
return 0;
static PUNICODE_STRING
psutil_threaded_get_filename(HANDLE hFile) {
DWORD dwWait;
HANDLE hThread;
DWORD result = 0;
ULONG attempts = 8;
psutil_query_thread threadContext = {
.status = STATUS_UNSUCCESSFUL,
.hFile = hFile,
.FileName = NULL,
.Length = 0x200,
};

threadContext.FileName = MALLOC_ZERO(threadContext.Length);
if (threadContext.FileName == NULL) {
PyErr_NoMemory();
return NULL;
}

// A loop is needed because the I/O subsystem likes to give us the
// wrong return lengths...
do {
status = NtQueryObject(
hFile,
ObjectNameInformation,
globalFileName,
bufferSize,
&bufferSize
);
if (status == STATUS_BUFFER_OVERFLOW ||
status == STATUS_INFO_LENGTH_MISMATCH ||
status == STATUS_BUFFER_TOO_SMALL)
hThread = CreateThread(
NULL, 0, (LPTHREAD_START_ROUTINE)psutil_get_filename,
&threadContext, 0, NULL);
if (hThread == NULL) {
PyErr_SetFromOSErrnoWithSyscall("CreateThread");
FREE(threadContext.FileName);
return NULL;
}

// Wait for the worker thread to finish.
dwWait = WaitForSingleObject(hThread, THREAD_TIMEOUT);

// If the thread hangs, kill it and cleanup.
if (dwWait == WAIT_TIMEOUT) {
psutil_debug(
"get handle name thread timed out after %i ms", THREAD_TIMEOUT);
if (TerminateThread(hThread, 0) == 0) {
PyErr_SetFromOSErrnoWithSyscall("TerminateThread");
result = 1;
}
} else if (dwWait == WAIT_FAILED) {
psutil_debug("WaitForSingleObject -> WAIT_FAILED");
result = 1;
if (TerminateThread(hThread, 0) == 0) {
PyErr_SetFromOSErrnoWithSyscall(
"WaitForSingleObject -> WAIT_FAILED -> TerminateThread");
} else {
PyErr_SetFromOSErrnoWithSyscall("WaitForSingleObject");
}
} else if (GetExitCodeThread(hThread, &result) == 0) {
result = 1;
PyErr_SetFromOSErrnoWithSyscall("GetExitCodeThread");
} else if (threadContext.status == STATUS_BUFFER_OVERFLOW ||
threadContext.status == STATUS_INFO_LENGTH_MISMATCH ||
threadContext.status == STATUS_BUFFER_TOO_SMALL)
{
FREE(globalFileName);
globalFileName = MALLOC_ZERO(bufferSize);
if (globalFileName == NULL) {
FREE(threadContext.FileName);
threadContext.FileName = MALLOC_ZERO(threadContext.Length);
if (threadContext.FileName == NULL) {
PyErr_NoMemory();
goto error;
CloseHandle(hThread);
return NULL;
}
}
else {
} else {
CloseHandle(hThread);
break;
}
} while (--attempts);

if (! NT_SUCCESS(status)) {
psutil_SetFromNTStatusErr(status, "NtQuerySystemInformation");
FREE(globalFileName);
globalFileName = NULL;
return 1;
}

return 0;

error:
if (globalFileName != NULL) {
FREE(globalFileName);
globalFileName = NULL;
}
return 1;
}


static DWORD
psutil_threaded_get_filename(HANDLE hFile) {
DWORD dwWait;
HANDLE hThread;
DWORD threadRetValue;

hThread = CreateThread(
NULL, 0, (LPTHREAD_START_ROUTINE)psutil_get_filename, &hFile, 0, NULL);
if (hThread == NULL) {
PyErr_SetFromOSErrnoWithSyscall("CreateThread");
return 1;
}
CloseHandle(hThread);

// Wait for the worker thread to finish.
dwWait = WaitForSingleObject(hThread, THREAD_TIMEOUT);
} while (--attempts && !result);

// If the thread hangs, kill it and cleanup.
if (dwWait == WAIT_TIMEOUT) {
psutil_debug(
"get handle name thread timed out after %i ms", THREAD_TIMEOUT);
if (TerminateThread(hThread, 0) == 0) {
PyErr_SetFromOSErrnoWithSyscall("TerminateThread");
CloseHandle(hThread);
return 1;
}
CloseHandle(hThread);
return 0;
if (result) {
FREE(threadContext.FileName);
return NULL;
}

if (dwWait == WAIT_FAILED) {
psutil_debug("WaitForSingleObject -> WAIT_FAILED");
if (TerminateThread(hThread, 0) == 0) {
PyErr_SetFromOSErrnoWithSyscall(
"WaitForSingleObject -> WAIT_FAILED -> TerminateThread");
CloseHandle(hThread);
return 1;
}
PyErr_SetFromOSErrnoWithSyscall("WaitForSingleObject");
CloseHandle(hThread);
return 1;
if (! NT_SUCCESS(threadContext.status)) {
psutil_SetFromNTStatusErr(threadContext.status, "NtQuerySystemInformation");
FREE(threadContext.FileName);
return NULL;
}

if (GetExitCodeThread(hThread, &threadRetValue) == 0) {
if (TerminateThread(hThread, 0) == 0) {
PyErr_SetFromOSErrnoWithSyscall(
"GetExitCodeThread (failed) -> TerminateThread");
CloseHandle(hThread);
return 1;
}

CloseHandle(hThread);
PyErr_SetFromOSErrnoWithSyscall("GetExitCodeThread");
return 1;
}
CloseHandle(hThread);
return threadRetValue;
return threadContext.FileName;
}


Expand All @@ -214,15 +197,12 @@ psutil_get_open_files(DWORD dwPid, HANDLE hProcess) {
ULONG i = 0;
BOOLEAN errorOccurred = FALSE;
PyObject* py_path = NULL;
PyObject* py_retlist = PyList_New(0);;
PyObject* py_retlist = PyList_New(0);
PUNICODE_STRING fileName = NULL;

if (!py_retlist)
return NULL;

// Due to the use of global variables, ensure only 1 call
// to psutil_get_open_files() is running.
EnterCriticalSection(&PSUTIL_CRITICAL_SECTION);

if (psutil_enum_handles(&handlesList) != 0)
goto error;

Expand All @@ -243,13 +223,20 @@ psutil_get_open_files(DWORD dwPid, HANDLE hProcess) {
continue;
}

// This will set *globalFileName* global variable.
if (psutil_threaded_get_filename(hFile) != 0)
if (GetFileType(hFile) != FILE_TYPE_DISK) {
SetLastError(0);
CloseHandle(hFile);
hFile = NULL;
continue;
}

fileName = psutil_threaded_get_filename(hFile);
if (fileName == NULL)
goto error;

if ((globalFileName != NULL) && (globalFileName->Length > 0)) {
py_path = PyUnicode_FromWideChar(globalFileName->Buffer,
wcslen(globalFileName->Buffer));
if ((fileName != NULL) && (fileName->Length > 0)) {
py_path = PyUnicode_FromWideChar(fileName->Buffer,
wcslen(fileName->Buffer));
if (! py_path)
goto error;
if (PyList_Append(py_retlist, py_path))
Expand All @@ -258,9 +245,9 @@ psutil_get_open_files(DWORD dwPid, HANDLE hProcess) {
}

// Loop cleanup section.
if (globalFileName != NULL) {
FREE(globalFileName);
globalFileName = NULL;
if (fileName != NULL) {
FREE(fileName);
fileName = NULL;
}
CloseHandle(hFile);
hFile = NULL;
Expand All @@ -276,16 +263,13 @@ psutil_get_open_files(DWORD dwPid, HANDLE hProcess) {
exit:
if (hFile != NULL)
CloseHandle(hFile);
if (globalFileName != NULL) {
FREE(globalFileName);
globalFileName = NULL;
}
if (fileName != NULL)
FREE(fileName);
if (py_path != NULL)
Py_DECREF(py_path);
if (handlesList != NULL)
FREE(handlesList);

LeaveCriticalSection(&PSUTIL_CRITICAL_SECTION);
if (errorOccurred == TRUE)
return NULL;
return py_retlist;
Expand Down

0 comments on commit b060339

Please sign in to comment.