From bf6577c8f4587c3f4de36658b2ae7a2f21a8bd3c Mon Sep 17 00:00:00 2001
From: Philipp Moritz <pcmoritz@gmail.com>
Date: Sun, 20 Dec 2020 12:10:28 -0800
Subject: [PATCH] Switch debugger to sockets and support unicode (#13004)

---
 python/ray/scripts/scripts.py         |  7 ++-----
 python/ray/tests/test_ray_debugger.py |  5 +++++
 python/ray/util/rpdb.py               | 25 +++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 7d72914ee..4a1dd6e28 100644
--- a/python/ray/scripts/scripts.py
+++ b/python/ray/scripts/scripts.py
@@ -6,7 +6,6 @@ import logging
 import os
 import subprocess
 import sys
-from telnetlib import Telnet
 import time
 import urllib
 import urllib.parse
@@ -172,8 +171,7 @@ def continue_debug_session():
                         ray.experimental.internal_kv._internal_kv_del(key)
                         return
                     host, port = session["pdb_address"].split(":")
-                    with Telnet(host, int(port)) as tn:
-                        tn.interact()
+                    ray.util.rpdb.connect_pdb_client(host, int(port))
                     ray.experimental.internal_kv._internal_kv_del(key)
                     continue_debug_session()
                     return
@@ -215,8 +213,7 @@ def debug(address):
                 ray.experimental.internal_kv._internal_kv_get(
                     active_sessions[index]))
             host, port = session["pdb_address"].split(":")
-            with Telnet(host, int(port)) as tn:
-                tn.interact()
+            ray.util.rpdb.connect_pdb_client(host, int(port))
 
 
 @cli.command()
diff --git a/python/ray/tests/test_ray_debugger.py b/python/ray/tests/test_ray_debugger.py
index adea19684..e271dd3f6 100644
--- a/python/ray/tests/test_ray_debugger.py
+++ b/python/ray/tests/test_ray_debugger.py
@@ -44,6 +44,7 @@ def test_ray_debugger_commands(shutdown_only):
 
     @ray.remote
     def f():
+        """We support unicode too: 🐛"""
         ray.util.pdb.set_trace()
 
     result1 = f.remote()
@@ -55,6 +56,10 @@ def test_ray_debugger_commands(shutdown_only):
     p.expect("Enter breakpoint index or press enter to refresh: ")
     p.sendline("0")
     p.expect("-> ray.util.pdb.set_trace()")
+    p.sendline("ll")
+    # Cannot use the 🐛 symbol here because pexpect doesn't support
+    # unicode, but this test also does nicely:
+    p.expect("unicode")
     p.sendline("c")
     p.expect("Enter breakpoint index or press enter to refresh: ")
     p.sendline("0")
diff --git a/python/ray/util/rpdb.py b/python/ray/util/rpdb.py
index 134e0b4ec..251dc25a4 100644
--- a/python/ray/util/rpdb.py
+++ b/python/ray/util/rpdb.py
@@ -8,6 +8,7 @@ import json
 import logging
 import os
 import re
+import select
 import socket
 import sys
 import uuid
@@ -234,3 +235,27 @@ def set_trace(breakpoint_uuid=None):
 def post_mortem():
     rdb = connect_ray_pdb(None, None, False, None)
     rdb.post_mortem()
+
+
+def connect_pdb_client(host, port):
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.connect((host, port))
+
+    while True:
+        # Get the list of sockets which are readable.
+        read_sockets, write_sockets, error_sockets = select.select(
+            [sys.stdin, s], [], [])
+
+        for sock in read_sockets:
+            if sock == s:
+                # Incoming message from remote debugger.
+                data = sock.recv(4096)
+                if not data:
+                    return
+                else:
+                    sys.stdout.write(data.decode())
+                    sys.stdout.flush()
+            else:
+                # User entered a message.
+                msg = sys.stdin.readline()
+                s.send(msg.encode())