package org.zalando.nakadi.service.subscription.state;

import org.zalando.nakadi.service.subscription.StreamingContext;

import javax.annotation.Nullable;
import java.io.IOException;

public class CleanupState extends State {
    private final Exception exception;

    public CleanupState(@Nullable final Exception e) {
        this.exception = e;
    }

    CleanupState() {
        this(null);
    }

    @Override
    public void onEnter() {
        try {
            getContext().unregisterAuthorizationUpdates();
        } catch (final RuntimeException ex) {
            getLog().error("Unexpected fail during removing callback for registration updates", ex);
        }
        try {
            if (null != exception) {
                getOut().onException(exception);
            }
        } finally {
            try {
                getContext().unregisterSession();

            } finally {
                switchState(StreamingContext.DEAD_STATE);
            }

            try {
                getContext().getZkClient().close();
            } catch (final IOException e) {
                getLog().error("Unexpected fail to release zk connection", e);
            }
            getContext().getCurrentSpan().finish();
        }
    }
}

@@ -889,32 +889,187 @@
     %7D finally %7B%0A
+                if (null != exception) %7B%0A                    getContext().switchStateImmediately(StreamingContext.DEAD_STATE);%0A                %7D else%7B%0A    
                 
@@ -1102,32 +1102,50 @@
 xt.DEAD_STATE);%0A
+                %7D%0A
             %7D%0A%0A 

Switch state immediately to DEAD state/do not schedule it