socket code fixes

2 files changed, 60 insertions, 29 deletions

diff --git a/module/remote/thriftbackend/Socket.py b/module/remote/thriftbackend/Socket.py
new file mode 100644
index 000000000..6ee850d07
--- /dev/null
+++ b/module/remote/thriftbackend/Socket.py
@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+
+import socket
+
+from thrift.transport.TSocket import TSocket, TServerSocket, TTransportException
+
+
+class Socket(TSocket):
+    def __init__(self, host='localhost', port=7228, ssl=False):
+        TSocket.__init__(self, host, port)
+        self.ssl = ssl
+
+    def open(self):
+        self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.handle.settimeout(self._timeout)
+        self.handle.connect((self.host, self.port))
+
+
+class ServerSocket(TServerSocket, Socket):
+    def __init__(self, port=7228, host="0.0.0.0", key="", cert=""):
+        self.host = host
+        self.port = port
+        self.handle = None
+
+    def listen(self):
+        self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        if hasattr(self.handle, 'set_timeout'):
+          self.handle.set_timeout(None)
+        self.handle.bind((self.host, self.port))
+        self.handle.listen(128)
\ No newline at end of file
diff --git a/module/remote/thriftbackend/ThriftTest.py b/module/remote/thriftbackend/ThriftTest.py
index 7fe790cd9..587ca184a 100644
--- a/module/remote/thriftbackend/ThriftTest.py
+++ b/module/remote/thriftbackend/ThriftTest.py
@@ -9,9 +9,9 @@ sys.path.append(path)
 
 from thriftgen.pyload import Pyload
 from thriftgen.pyload.ttypes import *
+from Socket import Socket
 
 from thrift import Thrift
-from thrift.transport import TSocket
 from thrift.transport import TTransport
 
 from Protocol import Protocol
@@ -47,42 +47,42 @@ bench(proxy.get_collector)
 print
 try:
 
-  # Make socket
-  transport = TSocket.TSocket('localhost', 7228)
+    # Make socket
+    transport = Socket('localhost', 7228)
 
-  # Buffering is critical. Raw sockets are very slow
-  transport = TTransport.TBufferedTransport(transport)
+    # Buffering is critical. Raw sockets are very slow
+    transport = TTransport.TBufferedTransport(transport)
 
-  # Wrap in a protocol
-  protocol = Protocol(transport)
+    # Wrap in a protocol
+    protocol = Protocol(transport)
 
-  # Create a client to use the protocol encoder
-  client = Pyload.Client(protocol)
+    # Create a client to use the protocol encoder
+    client = Pyload.Client(protocol)
 
-  # Connect!
-  transport.open()
-  
-  print "Login", client.login("User", "pyloadweb")
-  
-  bench(client.getServerVersion)
-  bench(client.statusServer)
-  bench(client.statusDownloads)
-  bench(client.getQueue)
-  bench(client.getCollector)
+    # Connect!
+    transport.open()
 
-  print 
-  print client.getServerVersion()
-  print client.statusServer()
-  print client.statusDownloads()
-  q =  client.getQueue()
+    print "Login", client.login("User", "pyloadweb")
 
-  for p in q:
+    bench(client.getServerVersion)
+    bench(client.statusServer)
+    bench(client.statusDownloads)
+    bench(client.getQueue)
+    bench(client.getCollector)
+
+    print
+    print client.getServerVersion()
+    print client.statusServer()
+    print client.statusDownloads()
+    q =  client.getQueue()
+
+    for p in q:
       data = client.getPackageData(p.pid)
       print data
       print "Package Name: ", data.name
 
-  # Close!
-  transport.close()
-  
+    # Close!
+    transport.close()
+
 except Thrift.TException, tx:
-  print 'ThriftExpection: %s' % (tx.message)
+    print 'ThriftExpection: %s' % (tx.message)