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
"""
19
Procedural generated graph with soft connection limit of 5, max limit of 20:
20
 - 5000 peers, 500MB: 18 iterations
21
 - 1000 peers, 10GB: p50 297 iterations, p100 384 iterations (65% ~ 84% speed)
22
"""
23
PEER_COUNT = 5000
24
PIECE_COUNT = 125
25
PIECE_TRANSMIT_LIMIT = 10  # Number of pieces uploaded/downloaded per iteration
26
SOFT_CONNECTION_LIMIT = 5
27
MAX_CONNECTION_LIMIT = 20
28

29

30
class Peer(object):
31
    def __init__(self, name, piece_count):
32
        self.name = name
33
        self.failed_connection_attempts = 0
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
        for n in range(PEER_COUNT):
94
            peer = Peer(str(n), PIECE_COUNT)
95
            if n > 0:
96
                random.shuffle(self.peers)
97
                for candidate in self.peers:
98
                    if len(candidate.neighbors) < MAX_CONNECTION_LIMIT:
99
                        peer.connect(candidate)
100
                        if len(peer.neighbors) >= SOFT_CONNECTION_LIMIT:
101
                            break
102
                    else:
103
                        peer.failed_connection_attempts += 1
104
                        if peer.failed_connection_attempts > 50:
105
                            break
106

107
            self.peers.append(peer)
108

109
        self.peers.sort()
110
        for peer in self.peers:
111
            neighbors_str = ""
112
            for neighbor in peer.neighbors:
113
                neighbors_str = neighbors_str + neighbor.name + "; "
114
            print ('Peer %s failed %d connection attempts. Connected to peers %s' % (
115
                peer.name, peer.failed_connection_attempts, neighbors_str))
116

117
        # Set peer 0 to be the seeder.
118
        self.peers[0].pieces = [1]*PIECE_COUNT
119
        self.peers[0].completed = len(self.peers[0].pieces)
120

121
    def start(self):
122
        time = 0
123
        while True:
124
            print ('current time: %d.' % time)
125
            time += 1
126

127
            plan = []
128
            for p in self.peers:
129
                if not p.done():
130
                    for j in range(0, PIECE_TRANSMIT_LIMIT):
131
                        plan.append(p)
132
            random.shuffle(plan)
133
            for p in plan:
134
                p.fetch_step(time)
135

136
            for p in self.peers:
137
                p.fetch_cleanup()
138

139
            done = True
140
            for p in self.peers:
141
                if p.completed != len(p.pieces):
142
                    done = False
143

144
            if done:
145
                break
146

147
            if time > 1000:
148
                break
149

150
        print ('Done. Total time: %d.' % time)
151

152

153
def main():
154
    peer_manager = PeerManager()
155
    peer_manager.start()
156

157
if __name__== "__main__":
158
     main()
159