Commits

Ning Sun  committed 0bff663

add VenusHeader and logic to read a venus message from socket

  • Participants
  • Parent commits fcf46bc

Comments (0)

Files changed (2)

File venus/client.py

-# -*- Coding!: utf-8 -*-
+# --*-- Coding: utf-8 --*--
 
 import gevent
 import gevent.socket
 import gevent.event
 
+from venus import codec
+
 class VenusError(Exception): pass
 class ConnectError(VenusError): pass
 
             ## 3. Deliver response into self.trans[packet_id]
             ## 
             ## reconnect to server if exception occurs
-            pass
+            try:
+                header = codec.VenusHeader()
+                header_data = self.sock.recv(codec.VenusHeader.HEADER_SIZE)
+                header.decode(header_data)
+                message_class = codec.VENUS_MESSAGES[header.type]
+                if message_class is not  None:
+                    message = message_class(header)
+                    body = self.sock.recv(header.size - codec.VenusHeader.HEADER_SIZE)
+                    message.body = body
+
+                    self.trans[header.transaction_id].set(message)
+                else:
+                    pass ### do something when invalid message received                    
+            except:
+                pass ### do something when error in reading data from sock
 
     def send(self, msg):        
         packet_id = self.trans_id
             self.loop.kill()
         if self.sock:
             self.sock.close()            
+        self.trans.clear()
 
 class Config(object):
     def __init__(self, **kwargs):

File venus/codec.py

 
 import struct
 
+class VenusHeader(object):
+    HEADER_SIZE = 24
+    def __init__(self):
+        self.size = 0
+        self.codec = ">ihihil"
+        self.protocol_version = 0
+        self.type = 0
+        self.content_type = 0
+        self.request_id = 0
+        self.transaction_id = 0
+
+    def decode(self, data):
+        self.size,
+        self.protocol_version,
+        self.type,
+        self.content_type,
+        self.request_id,
+        self.transaction_id = struct.unpack(self.codec)
+
+    def encode(self):
+        return struct.pack(self.codec, self.size, self.protocol_version, self.type,
+                           self.content_type, self.request_id, self.transaction_id)
+
 class VenusMessage(object):
-    def __init__(self):
-        self.header_size = 24
-        self.header_codec = ">ihihil"
+    def __init__(self, header):
+        self.header = header or VenusHeader()
         self.body_size = 0
         self.body = ""
 
     def set_trans_id(self, trans_id):
-        self.trans_id = trans_id
+        self.header.transaction_id = trans_id
 
     def set_protocol_version(self, version):
-        self.version = version
+        self.header.protocol_version = version
 
     def set_content_type(self, content_type):
-        self.content_type = content_type
+        self.header.content_type = content_type
 
     def set_request_id(self, request_id):
-        self.request_id = request_id
+        self.header.request_id = request_id
+
+    def encode(self):
+        return self.header.encode() + self.body
         
-    def write(self, fobj):
-        message_size = self.header_size + self.body_size
-        header = struct.pack(self.header_codec, message_size, self.version,
-                             self.type, self.content_type, self.request_id, self.trans_id)
-        fobj.write(header+body)
-        
-    def read(self, fobj):
-        message_size_bytes = fobj.read(4)
-        message_size = struct.unpack(">i", message_size_bytes)
-        message_header_and_body = fobj.read(message_size - 4)
-        message_header, message_body = message_header_and_body[:20], message_header_and_body[20:]
-        self.version, self.type, self.content_type, self.request_id, self.trans_id = struct.unpack(">hihil")
-        self.body = message_body
-
 VENUS_MESSAGES = {}
 def message_type(type_code):
     def wrapper(the_class):
     def set_body(self, body):
         self.body = body
         self.body_size = len(body)
+        self.header.size = VenusHeader.HEADER_SIZE + self.body_size
 
 @message_type(0x02000002)
 class VenusResponse(VenusMessage):