Gossip

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from socketserver import BaseRequestHandler, UDPServer
from socket import socket, AF_INET, SOCK_DGRAM

import time
import threading
import random

import sys

if len(sys.argv) != 2 or sys.argv[1] not in [
        '10001', '10002', '10003', '10004', '10005'
]:
    print('usage: ' + sys.argv[0] + ' 10001|10002|10003|10004|10005')
    sys.exit(1)

CURRENT_SERV_HOST = 'localhost'
CURRENT_SERV_PORT = int(sys.argv[1])
CURRENT_DATA_FILE = 'data_localhost_' + sys.argv[1]
### 下面是data_localhost_10001的例子：
### #key;value;version(or timestamp)
### a;100;3
### b;200;16
### c;300;4

PEERS = [
    "localhost:10001", "localhost:10002", "localhost:10003", "localhost:10004",
    "localhost:10005"
]

GOSSIP_INTERVAL = 3

MSG_TYPE_REQ = 'request'
MSG_TYPE_REP = 'replay'


def debug(s):
    print("[debug] " + s)


def get_peers(num):
    "Get random num peers (exclude itself) from PEERS"
    current_peer = CURRENT_SERV_HOST + ':' + str(CURRENT_SERV_PORT)
    return random.sample([peer for peer in PEERS if peer != current_peer], num)


def prepare_msg(msg_type):
    "第一行是消息的类型，后面每行都表示一个key,value,version三元组"
    msg = "msg_type=" + msg_type + '\n'
    try:
        with open(CURRENT_DATA_FILE, 'rt') as f:
            for line in f:
                if line.startswith('#'):
                    pass
                else:
                    msg += line
    except FileNotFoundError:
        pass
    return msg.encode('utf-8')


def send_msg(peers, msg):
    s = socket(AF_INET, SOCK_DGRAM)
    for _, peer in enumerate(peers):
        debug('send msg to ' + peer)
        debug('msg is:[' + msg.decode("utf-8") + ']')
        hostaddr = peer.split(':')[0]
        port = peer.split(':')[1]
        s.sendto(msg, (hostaddr, int(port)))


def active_thread_impl(gossip_interval):
    while True:
        time.sleep(gossip_interval)
        peers = get_peers(3)
        msg = prepare_msg(MSG_TYPE_REQ)
        send_msg(peers, msg)


def update_data(received_msg):
    data = {}
    # read data from file
    try:
        with open(CURRENT_DATA_FILE, 'rt') as f:
            for line in f:
                if not line or line.startswith('#'):
                    # empty or comment
                    pass
                else:
                    # format is: key;value;version(timestamp)
                    token = line.split(';', 3)
                    data[token[0]] = (token[1], token[2])
    except FileNotFoundError:
        pass

    need_update = False
    for _, line in enumerate(received_msg.split('\n')):
        if not line:  # empty line
            continue
        # format is: key;value;version(timestamp)
        # debug('received line=[' + line + ']')
        token = line.split(';', 3)
        key = token[0]
        # value = token[1]
        ver = token[2]
        # debug('rec version=[' + version + ']')
        if key in data:
            # update if received data is newer than current data
            ver_cur_data = data[key][1]
            if int(ver) > int(ver_cur_data):
                need_update = True
                data[key] = (token[1], token[2])
        else:
            need_update = True
            data[key] = (token[1], token[2])

    # write data back to file
    if need_update:
        merged_data = ''
        for key, value_and_ver in data.items():
            merged_data += key + ";" + value_and_ver[0].rstrip(
            ) + ";" + value_and_ver[1].rstrip() + '\n'
            # debug('merged_data=[' + merged_data + ']')
        with open(CURRENT_DATA_FILE, 'wt') as f:
            f.write(merged_data)


class MyRequestHandler(BaseRequestHandler):
    def handle(self):
        # print('Got connection from', self.client_address)
        # Get message and client socket
        buf, sock = self.request
        msg = buf.decode("utf-8")
        debug('Got message:[' + msg + ']')
        first_line = msg.split('\n')[0]
        if not first_line.startswith('msg_type='):
            debug('Got a bad message, just ignore it.')
            return
        msg_type = first_line[len('msg_type='):]
        # Send back replay msg if receive request msg
        if msg_type == MSG_TYPE_REQ:
            debug('Got a request msg, will send a replay msg.')
            # send replay msg
            my_msg = prepare_msg(MSG_TYPE_REQ)
            sock.sendto(my_msg, self.client_address)
        # update current peer's data
        update_data(msg[len(first_line + '\n'):])


def passive_thread_impl():
    serv = UDPServer((CURRENT_SERV_HOST, CURRENT_SERV_PORT), MyRequestHandler)
    serv.serve_forever()     # 一般采用UDP来实现Gossip数据交换


if __name__ == '__main__':
    t1 = threading.Thread(target=passive_thread_impl)
    t1.start()

    t2 = threading.Thread(target=active_thread_impl, args=(GOSSIP_INTERVAL, ))
    t2.start()