diff --git a/porcupine/plugins/run/__init__.py b/porcupine/plugins/run/__init__.py index 0e932658a..5274b6675 100644 --- a/porcupine/plugins/run/__init__.py +++ b/porcupine/plugins/run/__init__.py @@ -12,6 +12,7 @@ from porcupine import get_main_window, get_tab_manager, menubar, tabs, utils from porcupine.plugins import python_venv +from porcupine.settings import global_settings from . import common, dialog, history, no_terminal, terminal @@ -76,6 +77,11 @@ def on_new_filetab(tab: tabs.FileTab) -> None: def setup() -> None: + # Memory limit affects all commands, so that you can "set it and forget it". + # If you choose to run a different command, it won't reset annoyingly. + global_settings.add_option("run_mem_limit_enabled", default=False) + global_settings.add_option("run_mem_limit_value", default=1000 * 1000 * 1000) + get_tab_manager().add_filetab_callback(on_new_filetab) menubar.add_filetab_command( diff --git a/porcupine/plugins/run/common.py b/porcupine/plugins/run/common.py index 38c91db3e..d34082615 100644 --- a/porcupine/plugins/run/common.py +++ b/porcupine/plugins/run/common.py @@ -2,10 +2,13 @@ import dataclasses import os +import re +import sys from pathlib import Path -from typing import Dict, List, Optional +from typing import Callable, Dict, List, Optional from porcupine import tabs, utils +from porcupine.settings import global_settings @dataclasses.dataclass @@ -34,7 +37,7 @@ class ExampleCommand: class Context: - def __init__(self, tab: tabs.FileTab, key_id: int): + def __init__(self, tab: tabs.FileTab, key_id: int) -> None: assert tab.path is not None self.file_path = tab.path self.project_path = utils.find_project_root(tab.path) @@ -69,3 +72,76 @@ def prepare_env() -> dict[str, str]: ) return env + + +def mem_limit_to_string(limit: int) -> str: + if limit >= 1000 * 1000 * 1000: + number = limit / (1000 * 1000 * 1000) + suffix = "GB" + elif limit >= 1000 * 1000: + number = limit / (1000 * 1000) + suffix = "MB" + elif limit >= 1000: + number = limit / 1000 + suffix = "KB" + else: + number = limit + suffix = "B" + + # Show 2GB instead of 2.0GB + if number == int(number): + number = int(number) + + return str(number) + suffix + + +def string_to_mem_limit(string: str) -> int | None: + match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)([KMG]?)B?", string.replace(" ", "").upper()) + if match is None: + return None + + number_as_string, suffix = match.groups() + number = float(number_as_string) + if suffix == "": + limit = round(number) + elif suffix == "K": + limit = round(1000 * number) + elif suffix == "M": + limit = round(1000 * 1000 * number) + elif suffix == "G": + limit = round(1000 * 1000 * 1000 * number) + else: + raise NotImplementedError + + # Ban limits smaller than 10MB. Python needs about 17MB to start. + if limit < 10 * 1000 * 1000: + return None + return limit + + +# If statements must be outside the function definition because of mypy bugs +if sys.platform == "win32": + + def create_memory_limit_callback() -> Callable[[], None]: + raise NotImplementedError + +else: + import resource + + def create_memory_limit_callback() -> Callable[[], None]: + if not global_settings.get("run_mem_limit_enabled", bool): + return lambda: None + + limit = global_settings.get("run_mem_limit_value", int) + limit_string = mem_limit_to_string(limit) + + def callback() -> None: + try: + resource.setrlimit(resource.RLIMIT_AS, (limit, limit)) + except Exception as e: + # Avoid using anything that acquires IO locks. + # Still not great, because the GIL could deadlock. See warning in docs. + message = f"Limiting memory usage to {limit_string} failed: {e}\r\n" + os.write(2, message.encode("utf-8")) + + return callback diff --git a/porcupine/plugins/run/dialog.py b/porcupine/plugins/run/dialog.py index af9fb325b..bf49143a6 100644 --- a/porcupine/plugins/run/dialog.py +++ b/porcupine/plugins/run/dialog.py @@ -6,6 +6,7 @@ from typing import Callable, TypeVar from porcupine import get_main_window, textutils, utils +from porcupine.settings import global_settings from . import common, history @@ -77,13 +78,13 @@ def __init__(self, ctx: common.Context): entry_area, text="Run this command:", get_substituted_value=(lambda: self.get_command().format_command()), - validated_callback=self.update_run_button, + validated_callback=self.update_disabled_states, ) self.cwd = _FormattingEntryAndLabels( entry_area, text="In this directory:", get_substituted_value=(lambda: str(self.get_command().format_cwd())), - validated_callback=self.update_run_button, + validated_callback=self.update_disabled_states, ) ttk.Label(content_frame, text="Substitutions:").pack(anchor="w") @@ -120,11 +121,33 @@ def __init__(self, ctx: common.Context): self.window.bind("", (lambda e: self.terminal_var.set(False)), add=True) self.window.bind("", (lambda e: self.terminal_var.set(True)), add=True) + self._mem_limit_enable_var = tkinter.BooleanVar( + value=global_settings.get("run_mem_limit_enabled", bool) + ) + self._mem_limit_value_var = tkinter.StringVar( + value=common.mem_limit_to_string(global_settings.get("run_mem_limit_value", int)) + ) + + # Do not show memory usage options on Windows. + # Limiting the memory on Windows isn't supported. + memory_frame = ttk.Frame(content_frame) + if sys.platform != "win32": + memory_frame.pack(fill="x", pady=5) + + ttk.Checkbutton( + memory_frame, variable=self._mem_limit_enable_var, text="Limit memory usage to: " + ).pack(side="left") + self._mem_entry = ttk.Entry(memory_frame, textvariable=self._mem_limit_value_var, width=20) + self._mem_entry.pack(side="left") + + self._mem_limit_enable_var.trace_add("write", self.update_disabled_states) + self._mem_limit_value_var.trace_add("write", self.update_disabled_states) + self._repeat_bindings = [ utils.get_binding(f"<>") for key_id in range(4) ] self._repeat_var = tkinter.StringVar(value=self._repeat_bindings[ctx.key_id]) - self._repeat_var.trace_add("write", self.update_run_button) + self._repeat_var.trace_add("write", self.update_disabled_states) repeat_frame = ttk.Frame(content_frame) repeat_frame.pack(fill="x", pady=10) @@ -145,7 +168,7 @@ def __init__(self, ctx: common.Context): self.run_button.pack(side="left", fill="x", expand=True, padx=(5, 0)) self.run_clicked = False - for entry in [self.command.entry, self.cwd.entry]: + for entry in [self.command.entry, self.cwd.entry, self._mem_entry]: entry.bind("", (lambda e: self.run_button.invoke()), add=True) entry.bind("", (lambda e: self.window.destroy()), add=True) @@ -174,16 +197,32 @@ def get_command(self) -> common.Command: substitutions=self._substitutions, ) + def save_memory_limit_setting(self) -> None: + global_settings.set("run_mem_limit_enabled", self._mem_limit_enable_var.get()) + value = common.string_to_mem_limit(self._mem_limit_value_var.get()) + if value is not None: + global_settings.set("run_mem_limit_value", value) + def get_key_id(self) -> int: return self._repeat_bindings.index(self._repeat_var.get()) - def command_and_cwd_are_valid(self) -> bool: + def _all_entries_are_valid(self) -> bool: try: command = self.get_command().format_command() cwd = self.get_command().format_cwd() except (ValueError, KeyError, IndexError): return False - return bool(command.strip()) and cwd.is_dir() and cwd.is_absolute() + + return ( + bool(command.strip()) + and cwd.is_dir() + and cwd.is_absolute() + and self._repeat_var.get() in self._repeat_bindings + and ( + common.string_to_mem_limit(self._mem_limit_value_var.get()) is not None + or not self._mem_limit_enable_var.get() + ) + ) def _select_command_autocompletion(self, command: common.Command, prefix: str) -> None: assert command.command_format.startswith(prefix) @@ -210,12 +249,17 @@ def _autocomplete(self, event: tkinter.Event[tkinter.Entry]) -> str | None: return None - def update_run_button(self, *junk: object) -> None: - if self.command_and_cwd_are_valid() and self._repeat_var.get() in self._repeat_bindings: + def update_disabled_states(self, *junk: object) -> None: + if self._all_entries_are_valid(): self.run_button.config(state="normal") else: self.run_button.config(state="disabled") + if self._mem_limit_enable_var.get(): + self._mem_entry.config(state="normal") + else: + self._mem_entry.config(state="disabled") + def on_run_clicked(self) -> None: self.run_clicked = True self.window.destroy() @@ -236,5 +280,6 @@ def ask_command(ctx: common.Context) -> tuple[common.Command, int] | None: asker.window.wait_window() if asker.run_clicked: + asker.save_memory_limit_setting() return (asker.get_command(), asker.get_key_id()) return None diff --git a/porcupine/plugins/run/history.py b/porcupine/plugins/run/history.py index 2eaf6eb66..12330d386 100644 --- a/porcupine/plugins/run/history.py +++ b/porcupine/plugins/run/history.py @@ -3,6 +3,7 @@ import copy import dataclasses import json +import logging import sys from pathlib import Path from typing import Optional @@ -13,6 +14,8 @@ from . import common +_log = logging.getLogger(__name__) + @dataclasses.dataclass class _HistoryItem: diff --git a/porcupine/plugins/run/no_terminal.py b/porcupine/plugins/run/no_terminal.py index 5d48aaa0f..f48fb4422 100644 --- a/porcupine/plugins/run/no_terminal.py +++ b/porcupine/plugins/run/no_terminal.py @@ -98,6 +98,10 @@ def running(self) -> bool: def _thread_target(self, command: str, env: dict[str, str]) -> None: self._queue.put(("info", command + "\n")) + popen_kwargs = utils.subprocess_kwargs.copy() + if sys.platform != "win32": + popen_kwargs["preexec_fn"] = common.create_memory_limit_callback() + try: self._shell_process = subprocess.Popen( command, @@ -106,7 +110,7 @@ def _thread_target(self, command: str, env: dict[str, str]) -> None: stderr=subprocess.STDOUT, shell=True, env=env, - **utils.subprocess_kwargs, + **popen_kwargs, ) except OSError as e: self._queue.put(("error", f"{type(e).__name__}: {e}\n")) diff --git a/porcupine/plugins/run/terminal.py b/porcupine/plugins/run/terminal.py index 8fd0ae054..d753bc1fa 100644 --- a/porcupine/plugins/run/terminal.py +++ b/porcupine/plugins/run/terminal.py @@ -34,6 +34,8 @@ def _run_in_windows_cmd(command: str, cwd: Path, env: dict[str, str]) -> None: # new command prompt i found and it works :) we need cmd # because start is built in to cmd (lol) real_command = ["cmd", "/c", "start"] + real_command + + # no memory limit, don't know how to get that to work on windows subprocess.Popen(real_command, env=env) @@ -53,7 +55,11 @@ def _run_in_macos_terminal_app(command: str, cwd: Path, env: dict[str, str]) -> ) os.chmod(file.name, 0o755) - subprocess.Popen(["open", "-a", "Terminal.app", file.name], env=env) + subprocess.Popen( + ["open", "-a", "Terminal.app", file.name], + env=env, + preexec_fn=common.create_memory_limit_callback(), + ) # the terminal might be still opening when we get here, that's why # the file deletes itself # right now the file removes itself before it runs the actual command so @@ -118,7 +124,11 @@ def _run_in_x11_like_terminal(command: str, cwd: Path, env: dict[str, str]) -> N real_command = [str(run_script), str(cwd), command] real_command.extend(map(str, command)) - subprocess.Popen([terminal, "-e", " ".join(map(shlex.quote, real_command))], env=env) + subprocess.Popen( + [terminal, "-e", " ".join(map(shlex.quote, real_command))], + env=env, + preexec_fn=common.create_memory_limit_callback(), + ) # this figures out which terminal to use every time the user wants to run diff --git a/tests/test_run_plugin.py b/tests/test_run_plugin.py index 5ccdf8062..f31d31fa7 100644 --- a/tests/test_run_plugin.py +++ b/tests/test_run_plugin.py @@ -1,5 +1,6 @@ import os import shutil +import struct import sys import time from tkinter import ttk @@ -8,6 +9,7 @@ from porcupine import get_main_window, get_tab_manager, utils from porcupine.plugins.run import common, dialog, history, no_terminal, terminal +from porcupine.settings import global_settings @pytest.fixture(autouse=True) @@ -310,6 +312,46 @@ def test_crlf_on_any_platform(tmp_path, wait_until): wait_until(lambda: "foo\nbar" in get_output()) +@pytest.fixture(autouse=True) +def restore_memory_limits(): + yield + global_settings.reset("run_mem_limit_enabled") + global_settings.reset("run_mem_limit_value") + + +@pytest.mark.skipif(struct.calcsize("P") != 8, reason="assumes 64-bit") +@pytest.mark.xfail( + sys.platform == "win32", strict=True, reason="memory limits don't work on windows" +) +def test_memory_limit(tmp_path, wait_until): + global_settings.set("run_mem_limit_enabled", True) + global_settings.set("run_mem_limit_value", 100 * 1000 * 1000) # 100MB + + # 1 array element = pointer = 8 bytes (on a 64-bit system) + # 15*1000*1000 array elements = 120MB + no_terminal.run_command(f'{utils.quote(sys.executable)} -c "[0] * (15*1000*1000)"', tmp_path) + wait_until(lambda: "The process failed" in get_output()) + assert "MemoryError" in get_output() + + global_settings.set("run_mem_limit_value", 200 * 1000 * 1000) # 200MB + + no_terminal.run_command(f'{utils.quote(sys.executable)} -c "[0] * (15*1000*1000)"', tmp_path) + wait_until(lambda: "The process completed successfully" in get_output()) + assert "MemoryError" not in get_output() + + +@pytest.mark.xfail( + sys.platform == "win32", strict=True, reason="memory limits don't work on windows" +) +def test_error_when_setting_memory_limit(wait_until, tmp_path): + global_settings.set("run_mem_limit_enabled", True) + global_settings.set("run_mem_limit_value", 10**20) # too huge for 64-bit int, will fail + + no_terminal.run_command("echo 123", tmp_path) + wait_until(lambda: "The process completed successfully" in get_output()) + assert "Limiting memory usage to 100000000000GB failed" in get_output() + + def test_changing_current_file(filetab, tmp_path, wait_until): filetab.textwidget.insert("end", 'with open("foo.py", "w") as f: f.write("lol")') filetab.save_as(tmp_path / "foo.py")