From 10f986ce06054b22b54fd38f4f79bd44c71d6251 Mon Sep 17 00:00:00 2001
From: Janne Valkealahti <janne.valkealahti@gmail.com>
Date: Sun, 26 May 2019 15:20:03 +0100
Subject: [PATCH] Polish AbstractState

- Polish previous change to AbstractState by using a shared function
  completionStateListenerSink to track completion sink as functionality
  between entries to submachine/regions should be similar.
- Mostly relates to #743
---
 .../statemachine/state/AbstractState.java     | 61 +++++++------------
 1 file changed, 22 insertions(+), 39 deletions(-)

diff --git a/spring-statemachine-core/src/main/java/org/springframework/statemachine/state/AbstractState.java b/spring-statemachine-core/src/main/java/org/springframework/statemachine/state/AbstractState.java
index f40ada5e..ba850f9b 100644
--- a/spring-statemachine-core/src/main/java/org/springframework/statemachine/state/AbstractState.java
+++ b/spring-statemachine-core/src/main/java/org/springframework/statemachine/state/AbstractState.java
@@ -252,26 +252,7 @@ public abstract class AbstractState<S, E> extends LifecycleObjectSupport impleme
 		return Mono.defer(() -> {
 			if (submachine != null) {
 				Disposable disposable = Mono.just(submachine)
-					.flatMap(submachine -> {
-						return Mono.<Void>create(sink -> {
-							final StateMachineListener<S, E> l = new StateMachineListenerAdapter<S, E>() {
-
-								@Override
-								public void stateContext(StateContext<S, E> stateContext) {
-									if (stateContext.getStage() == Stage.STATEMACHINE_STOP) {
-										if (stateContext.getStateMachine() == submachine && submachine.isComplete()) {
-											completionListeners.remove(this);
-											submachine.removeStateListener(this);
-											if (completionListeners.isEmpty()) {
-												sink.success();
-											}
-										}
-									}
-								}
-							};
-							submachine.addStateListener(l);
-						});
-					})
+					.flatMap(submachine -> completionStateListenerSink(submachine))
 					// TODO: REACTOR this is causing cancel which breaks some things
 					// .then(handleStateDoOnComplete(context))
 					.then(Mono.fromRunnable(() -> notifyStateOnComplete(context)))
@@ -280,25 +261,7 @@ public abstract class AbstractState<S, E> extends LifecycleObjectSupport impleme
 			} else if (!regions.isEmpty()) {
 				// TODO: REACTOR we should handle disposable
 				Flux.fromIterable(regions)
-					.flatMap(region -> {
-						return Mono.<Void>create(sink -> {
-							final StateMachineListener<S, E> l = new StateMachineListenerAdapter<S, E>() {
-
-								@Override
-								public void stateContext(StateContext<S, E> stateContext) {
-									if (stateContext.getStage() == Stage.STATEMACHINE_STOP) {
-										if (stateContext.getStateMachine() == region && region.isComplete()) {
-											completionListeners.remove(this);
-											region.removeStateListener(this);
-											sink.success();
-										}
-									}
-								}
-							};
-							completionListeners.add(l);
-							region.addStateListener(l);
-							});
-					})
+					.flatMap(region -> completionStateListenerSink(region))
 					.then(handleStateDoOnComplete(context))
 					.then(Mono.fromRunnable(() -> notifyStateOnComplete(context)))
 					.subscribe();
@@ -471,6 +434,26 @@ public abstract class AbstractState<S, E> extends LifecycleObjectSupport impleme
 		}
 	}
 
+	private Mono<Void> completionStateListenerSink(Region<S, E> region) {
+		return Mono.create(sink -> {
+			final StateMachineListener<S, E> listener = new StateMachineListenerAdapter<S, E>() {
+
+				@Override
+				public void stateContext(StateContext<S, E> stateContext) {
+					if (stateContext.getStage() == Stage.STATEMACHINE_STOP) {
+						if (stateContext.getStateMachine() == region && region.isComplete()) {
+							completionListeners.remove(this);
+							region.removeStateListener(this);
+							sink.success();
+						}
+					}
+				}
+			};
+			completionListeners.add(listener);
+			region.addStateListener(listener);
+			});
+	}
+
 	private void disposeDisposables() {
 		Disposable disposable;
 		while ((disposable = disposables.poll()) != null) {