Serialize installation step for add-on store

Installing add-ons must happen on the main thread. They must be installed one by one. Downloads may happen concurrently, moving off the main thread allows NVDA to stay responsive. Description of change: Use a ThreadPoolExecutor to enable downloading in the background. This will allow scaling up the number of threads as required. Provided an abstraction to make cancellation/tracking of complete simpler. Addons are installed (if download has completed) when the user tries to exit the dialog.
seanbudd · Jun 5, 2023 · May 24, 2022 · May 24, 2022 · Jun 8, 2022 · Jun 10, 2022
commit 17824e47929d228cb717d4c6159531c0cdb1c281
@@ -15,7 +15,11 @@
 
 AddonApiVersionT = Tuple[int, int, int]
 
-CURRENT: AddonApiVersionT = (buildVersion.version_year, buildVersion.version_major, buildVersion.version_minor)
+CURRENT: AddonApiVersionT = (
+	buildVersion.version_year,
+	buildVersion.version_major,
+	buildVersion.version_minor
+)
 BACK_COMPAT_TO: AddonApiVersionT = (2022, 1, 0)
 """
 As BACK_COMPAT_TO is incremented, the changed / removed parts / or reasoning should be added below.

diff --git a/source/addonStore/dataManager.py b/source/addonStore/dataManager.py
@@ -6,16 +6,23 @@
 import enum
 import json
 import os
-import shutil
+from concurrent.futures import (
+	Future,
+	ThreadPoolExecutor
+)
 from datetime import datetime, timedelta
 
 from logHandler import log
 import requests
 
 import addonAPIVersion
+import typing
 from typing import (
 	Optional,
 	List,
+	Dict,
+	Callable,
+	Tuple,
 )
 from .models import (
 	_createModelFromData,
@@ -46,17 +53,119 @@ class CachedAddonModel:
 	cachedAt: datetime
 
 
+class AddonFileDownloader:
+	OnCompleteT = Callable[[AddonDetailsModel, Optional[os.PathLike]], None]
+
+	def __init__(self, cacheDir: os.PathLike):
+		self._cacheDir = cacheDir
+		self.progress: Dict[AddonDetailsModel, int] = {}  # Number of chunks received
+		self._pending: Dict[Future, Tuple[AddonDetailsModel, AddonFileDownloader.OnCompleteT]] = {}
+		self.complete: Dict[AddonDetailsModel, os.PathLike] = {}  # Path to downloaded file
+		self._executor = ThreadPoolExecutor(
+			max_workers=1,
+			thread_name_prefix="AddonDownloader",
+		)
+
+	def download(
+			self,
+			addonData: AddonDetailsModel,
+			onComplete: OnCompleteT
+	):
+		self.progress[addonData] = 0
+		f: Future = self._executor.submit(
+			self._download, addonData,
+		)
+		self._pending[f] = addonData, onComplete
+		f.add_done_callback(self._done)
+
+	def _done(self, downloadAddonFuture: Future):
+		isCancelled = downloadAddonFuture not in self._pending
+		addonId = "CANCELLED" if isCancelled else self._pending[downloadAddonFuture][0].addonId
+		log.debug(f"Done called for {addonId}")
+		if isCancelled:
+			log.debug("Download was cancelled, not calling onComplete")
+			return
+		if not downloadAddonFuture.done() or downloadAddonFuture.cancelled():
+			log.error("Logic error with download in BG thread.")
+			return
+		if downloadAddonFuture.exception():
+			log.error(f"Unhandled exception in _download", exc_info=downloadAddonFuture.exception())
+			return
+		addonData, onComplete = self._pending.pop(downloadAddonFuture)
+		del self.progress[addonData]
+		cacheFilePath: Optional[os.PathLike] = downloadAddonFuture.result()
+		self.complete[addonData] = cacheFilePath
+		onComplete(addonData, cacheFilePath)
+
+	def cancelAll(self):
+		log.debug("Cancelling all")
+		for f in self._pending.keys():
+			f.cancel()
+		self._executor.shutdown(wait=False)
+		self._executor = None
+		self.progress.clear()
+		self._pending.clear()
+
+	def _download(self, addonData: AddonDetailsModel) -> Optional[os.PathLike]:
+		log.debug(f"starting download: {addonData.addonId}")
+		cacheFilePath = os.path.join(
+			self._cacheDir,
+			self._getCacheFilenameForAddon(addonData)
+		)
+		inProgressFilePath = cacheFilePath + ".download"
+		if addonData not in self.progress:
+			log.debug("the download was cancelled before it started.")
+			return None  # The download was cancelled
+		try:
+			with requests.get(addonData.addonURL, stream=True) as r:
+				with open(inProgressFilePath, 'wb') as fd:
+					for chunk in r.iter_content(chunk_size=128000):
+						log.debug(f"Chunk download: {addonData.addonId}")
+						fd.write(chunk)
+						if addonData in self.progress:
+							self.progress[addonData] += 1
+						else:
+							log.debug(f"Cancelled download: {addonData.addonId}")
+							return None  # The download was cancelled
+		except requests.exceptions.RequestException as e:
+			log.debugWarning(f"Unable to download addon file: {e}")
+			return None
+		except OSError as e:
+			log.debugWarning(f"Unable to save addon file ({inProgressFilePath}): {e}")
+			return None
+		log.debug(f"Download complete: {inProgressFilePath}")
+		if os.path.exists(cacheFilePath):
+			log.debug(f"Cache file already exists, deleting prior to rename: {cacheFilePath}")
+			os.remove(cacheFilePath)
+		os.rename(src=inProgressFilePath, dst=cacheFilePath)
+		log.debug(f"Cache file available: {cacheFilePath}")
+		return typing.cast(os.PathLike, cacheFilePath)
+
+	@staticmethod
+	def _getCacheFilenameForAddon(addonData: AddonDetailsModel) -> str:
+		return f"{addonData.addonId}-{addonData.versionName}.nvda-addon"
+
+	def __del__(self):
+		if self._executor is not None:
+			self._executor.shutdown(wait=False)
+			self._executor = None
+
+
 class DataManager:
 	_cacheFilename: str = "_cachedLatestAvailableAddons.json"
 	_cachePeriod = timedelta(hours=6)
 
-	def __init__(self, cacheDirLocation: os.PathLike):
+	def __init__(self, userConfigLocation: os.PathLike):
+		cacheDirLocation = os.path.join(userConfigLocation, "addonStore")
 		self._lang = "en"
 		self._preferredChannel = Channel.ALL
 		self._cacheFile = os.path.join(cacheDirLocation, DataManager._cacheFilename)
-		self._addonDownloadCacheDir = cacheDirLocation
+		self._addonDownloadCacheDir = os.path.join(cacheDirLocation, "_dl")
 		self._availableAddonCache: Optional[CachedAddonModel] = self._getCachedAddonData()
 
+	def getFileDownloader(self) -> AddonFileDownloader:
+		return AddonFileDownloader(self._addonDownloadCacheDir)
+
 	def _getLatestAvailableAddonsData(self) -> Optional[bytes]:
 		url = _getAddonStoreURL(self._preferredChannel, self._lang, _getCurrentApiVersionForURL())
 		try:
@@ -111,23 +220,3 @@ def getLatestAvailableAddons(self) -> List[AddonDetailsModel]:
 					cachedAt=fetchTime,
 				)
 		return self._availableAddonCache.availableAddons
-
-	def _getCacheFilenameForAddon(self, addonData: AddonDetailsModel) -> str:
-		return f"{addonData.addonId}-{addonData.versionName}.nvda-addon"
-
-	def downloadAddonFile(self, addonData: AddonDetailsModel):
-		localFilePath = os.path.join(
-			self._addonDownloadCacheDir,
-			self._getCacheFilenameForAddon(addonData)
-		)
-		try:
-			with requests.get(addonData.addonURL, stream=True) as r:
-				with open(localFilePath, 'wb') as f:
-					shutil.copyfileobj(r.raw, f)
-		except requests.exceptions.RequestException as e:
-			log.debugWarning(f"Unable to download addon file: {e}")
-			return None
-		except OSError as e:
-			log.debugWarning(f"Unable to save addon file ({localFilePath}): {e}")
-			return None
-		return localFilePath
@@ -848,6 +848,11 @@ def _terminate(module, name=None):
 	except:
 		log.exception("Error terminating %s" % name)
 
+
+def isMainThread() -> bool:
+	return threading.get_ident() == mainThreadId
+
+
 def requestPump():
 	"""Request a core pump.
 	This will perform any queued activity.
@@ -858,7 +863,7 @@ def requestPump():
 	if not _pump or _isPumpPending:
 		return
 	_isPumpPending = True
-	if threading.get_ident() == mainThreadId:
+	if isMainThread():
 		_pump.Start(PUMP_MAX_DELAY, True)
 		return
 	# This isn't the main thread. wx timers cannot be run outside the main thread.
@@ -882,7 +887,7 @@ def callLater(delay, callable, *args, **kwargs):
 		# If NVDA has not fully initialized yet, the wxApp may not be initialized.
 		# wx.CallLater and wx.CallAfter requires the wxApp to be initialized.
 		raise NVDANotInitializedError("Cannot schedule callable, wx.App is not initialized")
-	if threading.get_ident() == mainThreadId:
+	if isMainThread():
 		return wx.CallLater(delay, _callLaterExec, callable, args, kwargs)
 	else:
 		return wx.CallAfter(wx.CallLater,delay, _callLaterExec, callable, args, kwargs)

diff --git a/source/gui/addonStoreGui/viewModels.py b/source/gui/addonStoreGui/viewModels.py
@@ -4,13 +4,17 @@
 # See the file COPYING for more details.
 import typing
 import enum
+from os import (
+	PathLike,
+)
 from enum import (
 	Enum,
 )
 from typing import (
 	List,
 	Optional,
 	Dict,
+	Tuple,
 )
 import threading
 
@@ -47,6 +51,8 @@ class AvailableAddonStatus(DisplayStringEnum):
 	# noinspection PyArgumentList
 	DOWNLOAD_FAILED = enum.auto()
 	# noinspection PyArgumentList
+	DOWNLOAD_SUCCESS = enum.auto()
+	# noinspection PyArgumentList
 	INSTALLING = enum.auto()
 	# noinspection PyArgumentList
 	INSTALL_FAILED = enum.auto()
@@ -71,6 +77,8 @@ def _displayStringLabels(self) -> Dict[Enum, str]:
 			# Translators: Status for addons shown in the add-on store dialog
 			AvailableAddonStatus.DOWNLOAD_FAILED: _("Download failed"),
 			# Translators: Status for addons shown in the add-on store dialog
+			AvailableAddonStatus.DOWNLOAD_SUCCESS: _("Downloaded"),
+			# Translators: Status for addons shown in the add-on store dialog
 			AvailableAddonStatus.INSTALLING: _("Installing"),
 			# Translators: Status for addons shown in the add-on store dialog
 			AvailableAddonStatus.INSTALL_FAILED: _("Install failed"),
@@ -179,7 +187,7 @@ def __init__(
 	def _itemDataUpdated(self, addonListItemVM: AddonListItemVM):
 		addonId: str = addonListItemVM.Id
 		log.debug(f"status: {addonListItemVM.status}")
-		log.debug(f"equal instances: {addonListItemVM == self._addons[addonId]}")
+		assert addonListItemVM == self._addons[addonId], "Must be the same instance."
 		if addonId in self._addonsFilteredOrdered:
 			log.debug("calling")
 			index = self._addonsFilteredOrdered.index(addonId)
@@ -462,6 +470,9 @@ def __init__(self, dataManager: DataManager):
 		self._addons: List[AddonDetailsModel] = []
 		self._filteredStatuses = [None, AvailableAddonStatus.RUNNING]
 
+		self._downloader = dataManager.getFileDownloader()
+		self._pendingInstalls: List[Tuple[AddonListItemVM, PathLike]] = []
+
 		self.listVM: AddonListVM = AddonListVM(
 			addons=_createListItemVMs(self._addons, self._filteredStatuses)
 		)
@@ -493,47 +504,62 @@ def _makeActionsList(self):
 			AddonActionVM(
 				# Translators: Label for a button that installs the selected addon
 				displayName=_("Install"),
-				actionHandler=self.installAddon,
+				actionHandler=self.getAddon,
 				validCheck=self.isInstallActionValid,
 				listItemVM=selectedListItem
 			),
 		]
 
-	def _installAddonInBG(self, listItemVM: AddonListItemVM):
-		try:
-			for status in self._doInstallAddon(listItemVM):
-				log.debug(f"{listItemVM.Id} status: {status}")
-				listItemVM.status = status
-		except TranslatedError as e:
-			log.debugWarning(f"Error during installation of {listItemVM.Id}", exc_info=True)
-			# ensure calling on the main thread.
-			core.callLater(delay=0, callable=self.hasError.notify, error=e)
+	def getAddon(self, listItemVM: AddonListItemVM):
+		listItemVM.status = AvailableAddonStatus.DOWNLOADING
+		log.debug(f"{listItemVM.Id} status: {listItemVM.status}")
+		self._downloader.download(listItemVM.model, self._downloadComplete)
 
-	def installAddon(self, listItemVM: AddonListItemVM):
-		threading.Thread(
->
-			name="install addon",
-			args=[listItemVM],
-		).start()
+	def _downloadComplete(self, addonDetails: AddonDetailsModel, fileDownloaded: Optional[PathLike]):
+		listItemVM: Optional[AddonListItemVM] = self.listVM._addons[addonDetails.addonId]
+		if listItemVM is None:
+			log.error(f"No list item VM for addon with id: {addonDetails.addonId}")
+			return
 
-	def _doInstallAddon(self, addon: AddonListItemVM) -> typing.Generator[AvailableAddonStatus, None, None]:
-		yield AvailableAddonStatus.DOWNLOADING
-		fileDownloaded = self._dataManager.downloadAddonFile(addon.model)
 		if fileDownloaded is None:
-			yield AvailableAddonStatus.DOWNLOAD_FAILED
-			raise TranslatedError(
-				displayMessage=pgettext(
-					"addonStore",
-					"Unable to download addon."
-				)
+			listItemVM.status = AvailableAddonStatus.DOWNLOAD_FAILED
+			log.debugWarning(f"Error during download of {listItemVM.Id}", exc_info=True)
+			e = TranslatedError(
+				# Translators: A message shown when downloading an add-on fails
+				displayMessage=pgettext("addonStore", "Unable to download add-on.")
 			)
-		yield AvailableAddonStatus.INSTALLING
+			# ensure calling on the main thread.
+			core.callLater(delay=0, callable=self.hasError.notify, error=e)
+			return
+
+		listItemVM.status = AvailableAddonStatus.DOWNLOAD_SUCCESS
+		log.debug(f"Queuing add-on for install on dialog exit: {listItemVM.Id}")
+		# Add-ons can have "installTasks", which often call the GUI assuming they are on the main thread.
+		self._pendingInstalls.append((listItemVM, fileDownloaded))
+
+	def installPending(self):
+		if not core.isMainThread():
+			# Add-ons can have "installTasks", which often call the GUI assuming they are on the main thread.
+			log.error("installation must happen on main thread.")
+		for listItemVM, fileDownloaded in self._pendingInstalls:
+			self._doInstall(listItemVM, fileDownloaded)
+
+	def _doInstall(self, listItemVM: AddonListItemVM, fileDownloaded: PathLike):
+		if not core.isMainThread():
+			# Add-ons can have "installTasks", which often call the GUI assuming they are on the main thread.
+			log.error("installation must happen on main thread.")
+			return
+		listItemVM.status = AvailableAddonStatus.INSTALLING
+		log.debug(f"{listItemVM.Id} status: {listItemVM.status}")
 		try:
 			installAddon(fileDownloaded)
 		except TranslatedError as e:
-			yield AvailableAddonStatus.INSTALL_FAILED
-			raise e
-		yield AvailableAddonStatus.INSTALLED
+			listItemVM.status = AvailableAddonStatus.INSTALL_FAILED
+			log.debug(f"{listItemVM.Id} status: {listItemVM.status}")
+			self.hasError.notify(error=e)
+			return
+		listItemVM.status = AvailableAddonStatus.INSTALLED
+		log.debug(f"{listItemVM.Id} status: {listItemVM.status}")
 
 	def refresh(self):
 		threading.Thread(target=self._getAddonsInBG, name="getAddonData").start()
@@ -550,6 +576,11 @@ def _getAddonsInBG(self):
 		self.detailsVM.listItem = self.listVM.getSelection()
 		log.debug("completed refresh")
 
+	def cancelDownloads(self):
+		for a in self._downloader.progress.keys():
+			self.listVM._addons[a.addonId].status = AvailableAddonStatus.AVAILABLE
+		self._downloader.cancelAll()
+
 
 class TranslatedError(Exception):
 	def __init__(self, displayMessage: str):
@@ -603,13 +634,14 @@ def getPreviouslyInstalledAddonById(addonId: str) -> Optional[addonHandler.Addon
 	return None
 
 
-def installAddon(addonPath: str) -> None:
+def installAddon(addonPath: PathLike) -> None:
 	""" Installs the addon at path.
 	Any error messages / warnings are presented to the user via a GUI message box.
 	If attempting to install an addon that is pending removal, it will no longer be pending removal.
 	@note See also L{gui.addonGui.installAddon}
 	@raise TranslatedError on failure
 	"""
+	addonPath = typing.cast(str, addonPath)
 	bundle = getAddonBundleToInstallIfValid(addonPath)
 	prevAddon = getPreviouslyInstalledAddonById(addonId=bundle.name)