From ee055873e613d7bedb474b348b5d30a874c4f0ee Mon Sep 17 00:00:00 2001
From: InsanusMokrassar <ovsyannikov.alexey95@gmail.com>
Date: Wed, 28 Dec 2022 22:48:49 +0600
Subject: [PATCH] fixes and improvements

---
 .../BehaviourContextWithFSM.kt                | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/tgbotapi.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSM.kt b/tgbotapi.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSM.kt
index fc9ac35f7b..cbec86bf52 100644
--- a/tgbotapi.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSM.kt
+++ b/tgbotapi.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSM.kt
@@ -169,11 +169,22 @@ class DefaultBehaviourContextWithFSM<T : State>(
                 statesManager.endChain(state)
             }
         }
+
+        fun Job.enableRemoveOnCompletion(state: T) {
+            invokeOnCompletion {
+                launchSafelyWithoutExceptions {
+                    if (this@enableRemoveOnCompletion === statesJobs[state]) {
+                        statesJobs.remove(state)
+                    }
+                }
+            }
+        }
+
         statesManager.onStartChain.subscribeSafelyWithoutExceptions(this) {
             statesJobsMutex.withLock {
                 runCatchingSafely { statesJobs.remove(it) ?.cancel() }
 
-                statesJobs[it] = launch { statePerformer(it) }
+                statesJobs[it] = launch { statePerformer(it) }.apply { enableRemoveOnCompletion(it) }
             }
         }
         statesManager.onEndChain.subscribeSafelyWithoutExceptions(this) {
@@ -186,7 +197,7 @@ class DefaultBehaviourContextWithFSM<T : State>(
             statesJobsMutex.withLock {
                 runCatchingSafely { statesJobs.remove(old) ?.cancel() }
                 runCatchingSafely { statesJobs.remove(new) ?.cancel() }
-                statesJobs[new] = launch { statePerformer(new) }
+                statesJobs[new] = launch { statePerformer(new) }.apply { enableRemoveOnCompletion(it) }
             }
             if (old.context != new.context) {
                 updatesFlows.remove(old.context)
@@ -194,7 +205,11 @@ class DefaultBehaviourContextWithFSM<T : State>(
         }
 
         statesManager.getActiveStates().forEach {
-            launch { statePerformer(it) }
+            statesJobsMutex.withLock {
+                runCatchingSafely { statesJobs.remove(it) ?.cancel() }
+
+                statesJobs[it] = launch { statePerformer(it) }.apply { enableRemoveOnCompletion(it) }
+            }
         }
     }
     /**