OkHTTP disconnects blocking

Author: GuilhermeJ

lib/src/main/java/ua/naiksoftware/stomp/OkHttpConnectionProvider.java

@@ -4,10 +4,15 @@
 import android.support.annotation.Nullable;
 import android.util.Log;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
 
+import io.reactivex.Completable;
+import io.reactivex.CompletableEmitter;
+import io.reactivex.CompletableOnSubscribe;
 import okhttp3.Headers;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
@@ -27,18 +32,22 @@ class OkHttpConnectionProvider extends AbstractConnectionProvider {
     @Nullable
     private WebSocket openedSocked;
 
+    @Nullable
+    private CompletableWebSocketListener currentListener;
+
     OkHttpConnectionProvider(String uri, @Nullable Map<String, String> connectHttpHeaders, OkHttpClient okHttpClient) {
         super();
         mUri = uri;
         mConnectHttpHeaders = connectHttpHeaders != null ? connectHttpHeaders : new HashMap<>();
         mOkHttpClient = okHttpClient;
     }
 
-    @NonNull
     @Override
     public void rawDisconnect() {
         if (openedSocked != null) {
             openedSocked.close(1000, "");
+            if (currentListener != null)
+                currentListener.getCompletable().blockingAwait();
         }
     }
 
@@ -49,52 +58,10 @@ void createWebSocketConnection() {
 
         addConnectionHeadersToBuilder(requestBuilder, mConnectHttpHeaders);
 
-        openedSocked = mOkHttpClient.newWebSocket(requestBuilder.build(),
-                new WebSocketListener() {
-                    @Override
-                    public void onOpen(WebSocket webSocket, @NonNull Response response) {
-                        LifecycleEvent openEvent = new LifecycleEvent(LifecycleEvent.Type.OPENED);
-
-                        TreeMap<String, String> headersAsMap = headersAsMap(response);
-
-                        openEvent.setHandshakeResponseHeaders(headersAsMap);
-                        emitLifecycleEvent(openEvent);
-                    }
-
-                    @Override
-                    public void onMessage(WebSocket webSocket, String text) {
-                        if (text.equals("\n"))
-                            Log.d(tag, "RECEIVED HEARTBEAT");
-                        else
-                            emitMessage(text);
-                    }
-
-                    @Override
-                    public void onMessage(WebSocket webSocket, @NonNull ByteString bytes) {
-                        emitMessage(bytes.utf8());
-                    }
-
-                    @Override
-                    public void onClosed(WebSocket webSocket, int code, String reason) {
-                        openedSocked = null;
-                        emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.CLOSED));
-                    }
-
-                    @Override
-                    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
-                        // in OkHttp, a Failure is equivalent to a JWS-Error *and* a JWS-Close
-                        emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.ERROR, new Exception(t)));
-                        openedSocked = null;
-                        emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.CLOSED));
-                    }
-
-                    @Override
-                    public void onClosing(final WebSocket webSocket, final int code, final String reason) {
-                        webSocket.close(code, reason);
-                    }
-                }
-
-        );
+        currentListener = new CompletableWebSocketListener();
+
+        openedSocked = mOkHttpClient
+                .newWebSocket(requestBuilder.build(), currentListener);
     }
 
     @Override
@@ -123,4 +90,76 @@ private void addConnectionHeadersToBuilder(@NonNull Request.Builder requestBuild
             requestBuilder.addHeader(headerEntry.getKey(), headerEntry.getValue());
         }
     }
+
+    private class CompletableWebSocketListener extends WebSocketListener implements CompletableOnSubscribe {
+        private List<CompletableEmitter> mEmitters = new ArrayList<>();
+        private Completable mCompletable = Completable.create(this);
+
+        @Override
+        public void onOpen(WebSocket webSocket, @NonNull Response response) {
+            LifecycleEvent openEvent = new LifecycleEvent(LifecycleEvent.Type.OPENED);
+
+            TreeMap<String, String> headersAsMap = headersAsMap(response);
+
+            openEvent.setHandshakeResponseHeaders(headersAsMap);
+            emitLifecycleEvent(openEvent);
+        }
+
+        @Override
+        public void onMessage(WebSocket webSocket, String text) {
+            if (text.equals("\n"))
+                Log.d(tag, "RECEIVED HEARTBEAT");
+            else
+                emitMessage(text);
+        }
+
+        @Override
+        public void onMessage(WebSocket webSocket, @NonNull ByteString bytes) {
+            emitMessage(bytes.utf8());
+        }
+
+        @Override
+        public void onClosed(WebSocket webSocket, int code, String reason) {
+            openedSocked = null;
+            emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.CLOSED));
+            launchCloseEvent();
+        }
+
+        @Override
+        public void onFailure(WebSocket webSocket, Throwable t, Response response) {
+            // in OkHttp, a Failure is equivalent to a JWS-Error *and* a JWS-Close
+            emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.ERROR, new Exception(t)));
+            openedSocked = null;
+            emitLifecycleEvent(new LifecycleEvent(LifecycleEvent.Type.CLOSED));
+            launchErrorEvent(t);
+        }
+
+        @Override
+        public void onClosing(final WebSocket webSocket, final int code, final String reason) {
+            webSocket.close(code, reason);
+        }
+
+        @Override
+        public void subscribe(CompletableEmitter emitter) {
+            mEmitters.add(emitter);
+        }
+
+        private Completable getCompletable(){
+            return mCompletable;
+        }
+
+        private void launchCloseEvent(){
+            for (CompletableEmitter emitter : mEmitters) {
+                if(!emitter.isDisposed())
+                    emitter.onComplete();
+            }
+        }
+
+        private void launchErrorEvent(Throwable t){
+            for (CompletableEmitter emitter : mEmitters) {
+                if(!emitter.isDisposed())
+                    emitter.onError(t);
+            }
+        }
+    }
 }