plomlompom.com Git - plomrogue2-experiments/blob - server.py

   1 #!/usr/bin/env python3
   2
   3 import socketserver
   4 import plom_socket_io
   5
   6 # Avoid "Address already in use" errors.
   7 socketserver.TCPServer.allow_reuse_address = True
   8
   9
  10 def fib(n):
  11     """Calculate n-th Fibonacci number."""
  12     if n in (1, 2):
  13         return 1
  14     else:
  15         return fib(n-1) + fib(n-2)
  16
  17
  18 def handle_message(message):
  19     """Evaluate message for computing-heavy tasks to perform, yield result.
  20
  21     Accepts one command: FIB, followed by positive integers, all tokens
  22     separated by whitespace. Will calculate and return for each such integer n
  23     the n-th Fibonacci number. Uses multiprocessing to perform multiple such
  24     calculations in parallel. Yields a 'CALCULATING …' message before the
  25     calculation starts, and finally yields a message containing the results.
  26     (The 'CALCULATING …' message coming before the results message is currently
  27     the main reason this works as a generator function using yield.)
  28
  29     When no command can be read into the message, just yields a 'NO COMMAND
  30     UNDERSTOOD:', followed by the message.
  31     """
  32     tokens = message.split(' ')
  33     if tokens[0] == 'FIB':
  34         msg_fail_fib = 'MALFORMED FIB REQUEST'
  35         if len(tokens) < 2:
  36             yield msg_fail_fib
  37             return
  38         numbers = []
  39         fail = False
  40         for token in tokens[1:]:
  41             if token != '0' and token.isdigit():
  42                 numbers += [int(token)]
  43             elif token == '':
  44                 continue
  45             else:
  46                 yield msg_fail_fib
  47                 return
  48         yield 'CALCULATING …'
  49         reply = ''
  50         from multiprocessing import Pool
  51         with Pool(len(numbers)) as p:
  52             results = p.map(fib, numbers)
  53         reply = ' '.join([str(r) for r in results])
  54         yield reply
  55         return
  56     yield 'NO COMMAND UNDERSTOOD: %s' % message
  57
  58
  59 class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
  60     """Enables threading on TCP server for asynchronous IO handling."""
  61     pass
  62
  63
  64 class MyTCPHandler(socketserver.BaseRequestHandler):
  65
  66     def handle(self):
  67         """Loop recv for input, act on it, send reply.
  68
  69         If input is 'QUIT', send reply 'BYE' and end loop / connection.
  70         Otherwise, use handle_message.
  71         """
  72
  73         print('CONNECTION FROM:', str(self.client_address))
  74         for message in plom_socket_io.recv(self.request):
  75             if message is None:
  76                 print('RECEIVED MALFORMED MESSAGE')
  77                 plom_socket_io.send(self.request, 'bad message')
  78             elif 'QUIT' == message:
  79                 plom_socket_io.send(self.request, 'BYE')
  80                 break
  81             else:
  82                 print('RECEIVED MESSAGE:', message)
  83                 for reply in handle_message(message):
  84                     plom_socket_io.send(self.request, reply)
  85         print('CONNECTION CLOSED:', str(self.client_address))
  86         self.request.close()
  87
  88
  89 server = ThreadedTCPServer(('localhost', 5000), MyTCPHandler)
  90 try:
  91     server.serve_forever()
  92 except KeyboardInterrupt:
  93     pass
  94 finally:
  95     server.server_close()