Skip to content

Commit

Permalink
Refactor code to better handle single instance lock (#77)
Browse files Browse the repository at this point in the history 
* Refactor.

* Move require to the top.
  • Loading branch information
@robinjhuang
robinjhuang authored Oct 12, 2024
1 parent c7b6650 commit 40a1b38
Showing 1 changed file with 129 additions and 144 deletions.
273 changes: 129 additions & 144 deletions src/main.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,11 @@ import * as Sentry from '@sentry/electron/main';
import { updateElectronApp, UpdateSourceType } from 'update-electron-app';
import * as net from 'net';

log.initialize();
// Handle creating/removing shortcuts on Windows when installing/uninstalling.
// Run this as early in the main process as possible.
if (require('electron-squirrel-startup')) app.quit();

updateElectronApp({
updateSource: {
type: UpdateSourceType.StaticStorage,
Expand All @@ -23,16 +28,9 @@ updateElectronApp({
updateInterval: '2 hours',
});

log.initialize();
// Handle creating/removing shortcuts on Windows when installing/uninstalling.
import('electron-squirrel-startup').then((ess) => {
const { default: check } = ess;
if (check) {
app.quit();
}
});

const gotTheLock = app.requestSingleInstanceLock();
const windowsLocalAppData = path.join(app.getPath('home'), 'comfyui-electron');
log.info('Windows Local App Data directory: ', windowsLocalAppData);

if (!gotTheLock) {
app.quit();
Expand All @@ -46,14 +44,13 @@ if (!gotTheLock) {
mainWindow.focus();
}
});
}

app.isPackaged &&
Sentry.init({
dsn: SENTRY_URL_ENDPOINT,
autoSessionTracking: false,
app.isPackaged &&
Sentry.init({
dsn: SENTRY_URL_ENDPOINT,
autoSessionTracking: false,

/* //WIP gather and send log from main
/* //WIP gather and send log from main
beforeSend(event, hint) {
hint.attachments = [
{
Expand All @@ -64,14 +61,124 @@ app.isPackaged &&
];
return event;
}, */
integrations: [
Sentry.childProcessIntegration({
breadcrumbs: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
events: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
}),
],
integrations: [
Sentry.childProcessIntegration({
breadcrumbs: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
events: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
}),
],
});

// This method will be called when Electron has finished
// initialization and is ready to create browser windows.
// Some APIs can only be used after this event occurs.

app.on('ready', async () => {
log.info('App ready');

app.on('activate', () => {
// On OS X it's common to re-create a window in the app when the
// dock icon is clicked and there are no other windows open.
if (BrowserWindow.getAllWindows().length === 0) {
const { userResourcesPath } = getResourcesPaths();
createWindow(userResourcesPath);
}
});

const { userResourcesPath, appResourcesPath } = getResourcesPaths();
log.info(`userResourcesPath: ${userResourcesPath}`);
log.info(`appResourcesPath: ${appResourcesPath}`);

try {
dotenv.config({ path: path.join(appResourcesPath, 'ComfyUI', '.env') });
} catch {
// if no .env file, skip it
}

createDirIfNotExists(userResourcesPath);

try {
await createWindow(userResourcesPath);
mainWindow.on('close', () => {
mainWindow = null;
app.quit();
});
ipcMain.on(IPC_CHANNELS.RENDERER_READY, () => {
log.info('Received renderer-ready message!');
// Send all queued messages
while (messageQueue.length > 0) {
const message = messageQueue.shift();
log.info('Sending queued message ', message.channel);
mainWindow.webContents.send(message.channel, message.data);
}
});
port = await findAvailablePort(8000, 9999).catch((err) => {
log.error(`ERROR: Failed to find available port: ${err}`);
throw err;
});
sendProgressUpdate('Setting up comfy environment...');
createComfyDirectories(userResourcesPath);
const pythonRootPath = path.join(userResourcesPath, 'python');
const pythonInterpreterPath =
process.platform === 'win32'
? path.join(pythonRootPath, 'python.exe')
: path.join(pythonRootPath, 'bin', 'python');
sendProgressUpdate('Setting up Python Environment...');
await setupPythonEnvironment(pythonInterpreterPath, appResourcesPath, userResourcesPath);
sendProgressUpdate('Starting Comfy Server...');
await launchPythonServer(pythonInterpreterPath, appResourcesPath, userResourcesPath);
} catch (error) {
log.error(error);
sendProgressUpdate(
'Was not able to start ComfyUI. Please check the logs for more details. You can open it from the tray icon.'
);
}

ipcMain.on(IPC_CHANNELS.RESTART_APP, () => {
log.info('Received restart app message!');
restartApp();
});
});

app.on('before-quit', () => {
if (progressInterval) {
clearInterval(progressInterval);
}
});

app.on('before-quit', () => {
if (progressInterval) {
clearInterval(progressInterval);
}
});

app.on('before-quit', async () => {
try {
await killPythonServer();
} catch (error) {
// Server did NOT exit properly
log.error('Python server did not exit properly');
log.error(error);
}
app.exit();
});

app.on('quit', () => {
log.info('Quitting ComfyUI');
app.exit();
});

// Quit when all windows are closed, except on macOS. There, it's common
// for applications and their menu bar to stay active until the user quits
// explicitly with Cmd + Q.
app.on('window-all-closed', () => {
log.info('Window all closed');
if (process.platform !== 'darwin') {
app.quit();
}
});
}

function loadComfyIntoMainWindow() {
if (!mainWindow) {
log.error('Trying to load ComfyUI into main window but it is not ready yet.');
Expand Down Expand Up @@ -178,6 +285,7 @@ export const createWindow = async (userResourcesPath: string): Promise<BrowserWi
mainWindow.hide();
app.dock.hide();
}
mainWindow = null;
});

const menu = buildMenu();
Expand All @@ -186,11 +294,6 @@ export const createWindow = async (userResourcesPath: string): Promise<BrowserWi
return mainWindow;
};

// Server Heartbeat Listener Variable
async function serverHeartBeat(): Promise<boolean> {
return isComfyServerReady(host, port);
}

const isComfyServerReady = async (host: string, port: number): Promise<boolean> => {
const url = `http://${host}:${port}/queue`;

Expand Down Expand Up @@ -306,77 +409,8 @@ function getResourcesPaths() {
return { userResourcesPath, appResourcesPath };
}

// This method will be called when Electron has finished
// initialization and is ready to create browser windows.
// Some APIs can only be used after this event occurs.
const windowsLocalAppData = path.join(app.getPath('home'), 'comfyui-electron');
log.info('Windows Local App Data directory: ', windowsLocalAppData);
app.on('ready', async () => {
log.info('App ready');

const { userResourcesPath, appResourcesPath } = getResourcesPaths();
log.info(`userResourcesPath: ${userResourcesPath}`);
log.info(`appResourcesPath: ${appResourcesPath}`);

try {
dotenv.config({ path: path.join(appResourcesPath, 'ComfyUI', '.env') });
} catch {
// if no .env file, skip it
}

createDirIfNotExists(userResourcesPath);

try {
await createWindow(userResourcesPath);
mainWindow.on('close', () => {
mainWindow = null;
app.quit();
});
ipcMain.on(IPC_CHANNELS.RENDERER_READY, () => {
log.info('Received renderer-ready message!');
// Send all queued messages
while (messageQueue.length > 0) {
const message = messageQueue.shift();
log.info('Sending queued message ', message.channel);
mainWindow.webContents.send(message.channel, message.data);
}
});
port = await findAvailablePort(8000, 9999).catch((err) => {
log.error(`ERROR: Failed to find available port: ${err}`);
throw err;
});
sendProgressUpdate('Setting up comfy environment...');
createComfyDirectories(userResourcesPath);
const pythonRootPath = path.join(userResourcesPath, 'python');
const pythonInterpreterPath =
process.platform === 'win32'
? path.join(pythonRootPath, 'python.exe')
: path.join(pythonRootPath, 'bin', 'python');
sendProgressUpdate('Setting up Python Environment...');
await setupPythonEnvironment(pythonInterpreterPath, appResourcesPath, userResourcesPath);
sendProgressUpdate('Starting Comfy Server...');
await launchPythonServer(pythonInterpreterPath, appResourcesPath, userResourcesPath);
} catch (error) {
log.error(error);
sendProgressUpdate(
'Was not able to start ComfyUI. Please check the logs for more details. You can open it from the tray icon.'
);
}

ipcMain.on(IPC_CHANNELS.RESTART_APP, () => {
log.info('Received restart app message!');
restartApp();
});
});

/** Interval to send progress updates to the renderer. */
let progressInterval: NodeJS.Timeout | null = null;
interface ProgressOptions {
endPercentage?: number;
duration?: number;
steps?: number;
overwrite?: boolean;
}

function sendProgressUpdate(status: string): void {
if (mainWindow) {
Expand Down Expand Up @@ -408,18 +442,6 @@ function sendProgressUpdate(status: string): void {
}
}

app.on('before-quit', () => {
if (progressInterval) {
clearInterval(progressInterval);
}
});

app.on('before-quit', () => {
if (progressInterval) {
clearInterval(progressInterval);
}
});

const killPythonServer = async (): Promise<void> => {
if (pythonProcess) {
log.info('Killing python server.');
Expand Down Expand Up @@ -448,43 +470,6 @@ const killPythonServer = async (): Promise<void> => {
}
};

app.on('before-quit', async () => {
try {
await killPythonServer();
} catch (error) {
// Server did NOT exit properly
log.error('Python server did not exit properly');
log.error(error);
}
app.exit();
});

app.on('quit', () => {
log.info('Quitting ComfyUI');
app.exit();
});

// Quit when all windows are closed, except on macOS. There, it's common
// for applications and their menu bar to stay active until the user quits
// explicitly with Cmd + Q.
app.on('window-all-closed', () => {
log.info('Window all closed');
if (process.platform !== 'darwin') {
app.quit();
}
});

app.on('activate', () => {
// On OS X it's common to re-create a window in the app when the
// dock icon is clicked and there are no other windows open.
if (process.platform === 'darwin') {
if (BrowserWindow.getAllWindows().length === 0) {
const { userResourcesPath } = getResourcesPaths();
createWindow(userResourcesPath);
}
}
});

const spawnPython = (
pythonInterpreterPath: string,
cmd: string[],
Expand Down

0 comments on commit 40a1b38

Please sign in to comment.