kraken

1
# Copyright (c) 2016-2019 Uber Technologies, Inc.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14
import json
15
import random
16
import sets
17

18
import networkx as nx
19

20
"""
21
Random regular graph with connection limit of 5:
22
 - 5000 peers, 500MB: 17 iterations
23
 - 1000 peers, 10GB: p50 294 iterations, p100 298 iterations (84% ~ 85% speed)
24
"""
25
PEER_COUNT = 5000
26
PIECE_COUNT = 125
27
PIECE_TRANSMIT_LIMIT = 10  # Number of pieces uploaded/downloaded per iteration
28
DEGREE = 5
29

30

31
class Peer(object):
32
    def __init__(self, name, piece_count):
33
        self.name = name
34
        self.neighbors = sets.Set()
35
        self.pieces = [0]*piece_count
36
        self.completed = 0
37
        self.time = 0
38

39
        self.uploaded_current_turn = 0
40
        self.downloaded_current_turn = 0
41

42
    def connect(self, other):
43
        self.neighbors.add(other)
44
        other.neighbors.add(self)
45

46
    def done(self):
47
        return self.completed == len(self.pieces)
48

49
    def fetch_step(self, time):
50
        if self.done():
51
            return
52

53
        if self.downloaded_current_turn >= PIECE_TRANSMIT_LIMIT:
54
            return
55

56
        candidates = []
57
        for n in self.neighbors:
58
            if n.uploaded_current_turn >= PIECE_TRANSMIT_LIMIT:
59
                continue
60

61
            for i in range(0, len(self.pieces)):
62
                if n.uploaded_current_turn >= PIECE_TRANSMIT_LIMIT:
63
                    continue
64

65
                if n.pieces[i] == 1 and self.pieces[i] == 0:
66
                    candidates.append((n, i))
67

68
        if len(candidates) == 0:
69
            return
70

71
        c = random.choice(candidates)
72

73
        self.pieces[c[1]] = 1
74
        self.completed += 1
75
        self.downloaded_current_turn += 1
76
        c[0].uploaded_current_turn += 1
77

78
        print ('Peer %s downloaded one piece from neighbor %s. Total completed: %d.' % (self.name, c[0].name, self.completed))
79

80
        if self.completed == len(self.pieces)-1:
81
            self.time = time
82
            print ('Peer %s finished downloading at time %d.' % (self.name, time))
83

84
    def fetch_cleanup(self):
85
        self.uploaded_current_turn = 0
86
        self.downloaded_current_turn = 0
87

88
class PeerManager(object):
89

90
    def __init__(self):
91
        self.peers = []
92

93
        g = nx.random_regular_graph(DEGREE, PEER_COUNT)
94
        for n in g:
95
            peer = Peer(str(n), PIECE_COUNT)
96
            self.peers.append(peer)
97

98
        for e in g.edges():
99
            self.peers[e[0]].connect(self.peers[e[1]])
100

101
        for peer in self.peers:
102
            neighbors_str = ""
103
            for neighbor in peer.neighbors:
104
                neighbors_str = neighbors_str + neighbor.name + "; "
105
            print ('Peer %s is connected to peers %s' % (peer.name, neighbors_str))
106

107
        # Set peer 0 to be the seeder.
108
        self.peers[0].pieces = [1]*PIECE_COUNT
109
        self.peers[0].completed = len(self.peers[0].pieces)
110

111
    def start(self):
112
        time = 0
113
        while True:
114
            print ('current time: %d.' % time)
115
            time += 1
116

117
            plan = []
118
            for p in self.peers:
119
                if not p.done():
120
                    for j in range(0, PIECE_TRANSMIT_LIMIT):
121
                        plan.append(p)
122
            random.shuffle(plan)
123
            for p in plan:
124
                p.fetch_step(time)
125

126
            for p in self.peers:
127
                p.fetch_cleanup()
128

129
            done = True
130
            for p in self.peers:
131
                if p.completed != len(p.pieces):
132
                    done = False
133

134
            if done:
135
                break
136

137
            if time > 1000:
138
                break
139

140
        print ('Done. Total time: %d.' % time)
141

142

143
def main():
144
    peer_manager = PeerManager()
145
    peer_manager.start()
146

147
if __name__== "__main__":
148
     main()
149