From 405079f744fb6901a97a1910218ee008a03607de Mon Sep 17 00:00:00 2001
From: Aanand Prasad <aanand.prasad@gmail.com>
Date: Mon, 20 Jan 2014 15:52:07 +0000
Subject: [PATCH] Use raw socket in 'fig run', simplify _attach_to_container

---
 fig/cli/main.py         | 41 ++++++++---------------------------------
 fig/cli/socketclient.py | 14 +++++++-------
 2 files changed, 15 insertions(+), 40 deletions(-)

diff --git a/fig/cli/main.py b/fig/cli/main.py
index 1696a7c98..66d2bdf28 100644
--- a/fig/cli/main.py
+++ b/fig/cli/main.py
@@ -220,12 +220,7 @@ class TopLevelCommand(Command):
             service.start_container(container, ports=None)
             print(container.name)
         else:
-            with self._attach_to_container(
-                container.id,
-                interactive=True,
-                logs=True,
-                raw=tty
-            ) as c:
+            with self._attach_to_container(container.id, raw=tty) as c:
                 service.start_container(container, ports=None)
                 c.run()
 
@@ -317,35 +312,15 @@ class TopLevelCommand(Command):
                 print("Gracefully stopping... (press Ctrl+C again to force)")
                 self.project.stop(service_names=options['SERVICE'])
 
-    def _attach_to_container(self, container_id, interactive, logs=False, stream=True, raw=False):
-        stdio = self.client.attach_socket(
-            container_id,
-            params={
-                'stdin': 1 if interactive else 0,
-                'stdout': 1,
-                'stderr': 0,
-                'logs': 1 if logs else 0,
-                'stream': 1 if stream else 0
-            },
-            ws=True,
-        )
-
-        stderr = self.client.attach_socket(
-            container_id,
-            params={
-                'stdin': 0,
-                'stdout': 0,
-                'stderr': 1,
-                'logs': 1 if logs else 0,
-                'stream': 1 if stream else 0
-            },
-            ws=True,
-        )
+    def _attach_to_container(self, container_id, raw=False):
+        socket_in = self.client.attach_socket(container_id, params={'stdin': 1, 'stream': 1})
+        socket_out = self.client.attach_socket(container_id, params={'stdout': 1, 'logs': 1, 'stream': 1})
+        socket_err = self.client.attach_socket(container_id, params={'stderr': 1, 'logs': 1, 'stream': 1})
 
         return SocketClient(
-            socket_in=stdio,
-            socket_out=stdio,
-            socket_err=stderr,
+            socket_in=socket_in,
+            socket_out=socket_out,
+            socket_err=socket_err,
             raw=raw,
         )
 
diff --git a/fig/cli/socketclient.py b/fig/cli/socketclient.py
index bbae60e6d..842c9c6a4 100644
--- a/fig/cli/socketclient.py
+++ b/fig/cli/socketclient.py
@@ -57,15 +57,15 @@ class SocketClient:
 
     def run(self):
         if self.socket_in is not None:
-            self.start_background_thread(target=self.send_ws, args=(self.socket_in, sys.stdin))
+            self.start_background_thread(target=self.send, args=(self.socket_in, sys.stdin))
 
         recv_threads = []
 
         if self.socket_out is not None:
-            recv_threads.append(self.start_background_thread(target=self.recv_ws, args=(self.socket_out, sys.stdout)))
+            recv_threads.append(self.start_background_thread(target=self.recv, args=(self.socket_out, sys.stdout)))
 
         if self.socket_err is not None:
-            recv_threads.append(self.start_background_thread(target=self.recv_ws, args=(self.socket_err, sys.stderr)))
+            recv_threads.append(self.start_background_thread(target=self.recv, args=(self.socket_err, sys.stderr)))
 
         for t in recv_threads:
             t.join()
@@ -76,10 +76,10 @@ class SocketClient:
         thread.start()
         return thread
 
-    def recv_ws(self, socket, stream):
+    def recv(self, socket, stream):
         try:
             while True:
-                chunk = socket.recv()
+                chunk = socket.recv(4096)
 
                 if chunk:
                     stream.write(chunk)
@@ -89,7 +89,7 @@ class SocketClient:
         except Exception as e:
             log.debug(e)
 
-    def send_ws(self, socket, stream):
+    def send(self, socket, stream):
         while True:
             r, w, e = select([stream.fileno()], [], [])
 
@@ -97,7 +97,7 @@ class SocketClient:
                 chunk = stream.read(1)
 
                 if chunk == '':
-                    socket.send_close()
+                    socket.close()
                     break
                 else:
                     try: