From a52e0c5b6a6bf32ea43889c65f07fa8a09d4b2b2 Mon Sep 17 00:00:00 2001 From: ShiftHackZ Date: Sat, 10 Aug 2024 10:39:25 +0300 Subject: [PATCH] Handle LocalDiffusion process death in TextToImageTask.kt --- .../data/preference/PreferenceManagerImpl.kt | 7 +++++++ .../domain/preference/PreferenceManager.kt | 1 + .../aisdv1/work/task/TextToImageTask.kt | 20 ++++++++++++++++++- 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/data/src/main/java/com/shifthackz/aisdv1/data/preference/PreferenceManagerImpl.kt b/data/src/main/java/com/shifthackz/aisdv1/data/preference/PreferenceManagerImpl.kt index c9020a61..224449af 100644 --- a/data/src/main/java/com/shifthackz/aisdv1/data/preference/PreferenceManagerImpl.kt +++ b/data/src/main/java/com/shifthackz/aisdv1/data/preference/PreferenceManagerImpl.kt @@ -212,6 +212,12 @@ class PreferenceManagerImpl( .apply() .also { onPreferencesChanged() } + override var backgroundProcessCount: Int + get() = preferences.getInt(KEY_BACKGROUND_PROCESS_COUNT, 0) + set(value) = preferences.edit() + .putInt(KEY_BACKGROUND_PROCESS_COUNT, value) + .apply() + override var galleryGrid: Grid get() = preferences.getInt(KEY_GALLERY_GRID, 0).let { Grid.entries[it] } set(value) = preferences.edit() @@ -273,6 +279,7 @@ class PreferenceManagerImpl( const val KEY_DESIGN_COLOR_TOKEN = "key_design_color_token_theme" const val KEY_DESIGN_DARK_TOKEN = "key_design_dark_color_token_theme" const val KEY_BACKGROUND_GENERATION = "key_background_generation" + const val KEY_BACKGROUND_PROCESS_COUNT = "key_background_process_count" const val KEY_GALLERY_GRID = "key_gallery_grid" } } diff --git a/domain/src/main/java/com/shifthackz/aisdv1/domain/preference/PreferenceManager.kt b/domain/src/main/java/com/shifthackz/aisdv1/domain/preference/PreferenceManager.kt index 34258396..bc52c431 100644 --- a/domain/src/main/java/com/shifthackz/aisdv1/domain/preference/PreferenceManager.kt +++ b/domain/src/main/java/com/shifthackz/aisdv1/domain/preference/PreferenceManager.kt @@ -32,6 +32,7 @@ interface PreferenceManager { var designColorToken: String var designDarkThemeToken: String var backgroundGeneration: Boolean + var backgroundProcessCount: Int var galleryGrid: Grid fun observe(): Flowable diff --git a/feature/work/src/main/java/com/shifthackz/aisdv1/work/task/TextToImageTask.kt b/feature/work/src/main/java/com/shifthackz/aisdv1/work/task/TextToImageTask.kt index 996f2e5a..e8a43efc 100644 --- a/feature/work/src/main/java/com/shifthackz/aisdv1/work/task/TextToImageTask.kt +++ b/feature/work/src/main/java/com/shifthackz/aisdv1/work/task/TextToImageTask.kt @@ -23,9 +23,9 @@ internal class TextToImageTask( workerParameters: WorkerParameters, pushNotificationManager: PushNotificationManager, activityIntentProvider: ActivityIntentProvider, - preferenceManager: PreferenceManager, observeHordeProcessStatusUseCase: ObserveHordeProcessStatusUseCase, observeLocalDiffusionProcessStatusUseCase: ObserveLocalDiffusionProcessStatusUseCase, + private val preferenceManager: PreferenceManager, private val backgroundWorkObserver: BackgroundWorkObserver, private val textToImageUseCase: TextToImageUseCase, private val fileProviderDescriptor: FileProviderDescriptor, @@ -45,12 +45,26 @@ internal class TextToImageTask( override val genericNotificationId: Int = NOTIFICATION_TEXT_TO_IMAGE_GENERIC override fun createWork(): Single { + // Workaround for LocalDiffusion provider: + // + // If LocalDiffusion process previously died, prevent WorkManager to go to infinite + // task repeat loop. + if (preferenceManager.backgroundProcessCount > 0) { + handleProcess() + handleError(Throwable("Background process count > 0")) + compositeDisposable.clear() + return Single.just(Result.failure()) + } + + preferenceManager.backgroundProcessCount++ handleStart() backgroundWorkObserver.refreshStatus() backgroundWorkObserver.dismissResult() + return try { val file = File(fileProviderDescriptor.workCacheDirPath, Constants.FILE_TEXT_TO_IMAGE) if (!file.exists()) { + preferenceManager.backgroundProcessCount-- handleError(Throwable("File is null.")) compositeDisposable.clear() return Single.just(Result.failure()) @@ -60,6 +74,7 @@ internal class TextToImageTask( val payload = bytes.toTextToImagePayload() if (payload == null) { + preferenceManager.backgroundProcessCount-- handleError(Throwable("Payload is null.")) compositeDisposable.clear() return Single.just(Result.failure()) @@ -71,15 +86,18 @@ internal class TextToImageTask( textToImageUseCase(payload) .doOnSubscribe { handleProcess() } .map { result -> + preferenceManager.backgroundProcessCount-- handleSuccess(result) Result.success() } .onErrorReturn { t -> + preferenceManager.backgroundProcessCount-- handleError(t) Result.failure() } .doFinally { compositeDisposable.clear() } } catch (e: Exception) { + preferenceManager.backgroundProcessCount-- handleError(e) compositeDisposable.clear() Single.just(Result.failure())