Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/peer.py: 87.18%

1# Copyright (C) 2019 Bloomberg LP 

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 

15import logging 

16from threading import Lock 

17 

18 

19class Peer: 

20 """Represents a client during a session.""" 

21 

22 # We will keep a global list of Peers indexed by their `Peer.uid`. 

23 __peers_by_uid = {} # type: ignore 

24 __peers_by_uid_lock = Lock() 

25 

26 @classmethod 

27 def register_peer(cls, uid, context, token=None): 

28 """Registers a new peer from a given context. 

29 

30 Args: 

31 uid (str): a unique identifier 

32 token (str): an authentication token (optional) 

33 context (grpc.ServicerContext): context in which the peer is 

34 being registered 

35 

36 Returns: 

37 Peer: an existing or newly created Peer object 

38 """ 

39 

40 with cls.__peers_by_uid_lock: 

41 # If the peer exists, we just increment the counter on the existing 

42 # instance: 

43 if uid in cls.__peers_by_uid: 

44 existing_peer = cls.__peers_by_uid[uid] 

45 logging.getLogger(__name__).debug('Registering another instance ' 

46 f'of Peer with uid {uid} ') 

47 existing_peer.__instance_count += 1 

48 return existing_peer 

49 else: 

50 # Otherwise we add ourselves to the list of Peers: 

51 new_peer = Peer(uid=uid, token=token) 

52 cls.__peers_by_uid[uid] = new_peer 

53 return cls.__peers_by_uid[uid] 

54 

55 def __init__(self, uid, token=None, tool_name=None, tool_version=None): 

56 """Creates a new Peer object. 

57 

58 Args: 

59 uid (str): a unique identifier 

60 token (str): an authentication token (optional) 

61 """ 

62 self._uid = uid # This uniquely identifies a client 

63 self._token = token 

64 

65 # Each Peer object contains the number of instances of itself: 

66 self.__instance_count = 1 

67 

68 @classmethod 

69 def find_peer(cls, uid): 

70 return cls.__peers_by_uid.get(uid, None) 

71 

72 def __eq__(self, other): 

73 if not isinstance(other, Peer): 

74 return False 

75 

76 return self.uid == other.uid and self.token == other.token 

77 

78 def __hash__(self): 

79 return hash(self.uid) # This string is unique for each peer 

80 

81 def __str__(self): 

82 return f'Peer: uid: {self._uid}' 

83 

84 @property 

85 def uid(self): 

86 return self._uid 

87 

88 @property 

89 def token(self): 

90 return self._token 

91 

92 @classmethod 

93 def deregister_peer(cls, peer_uid): 

94 """Deregisters a Peer from the list of peers present. 

95 If the Peer deregistered has a single instance, we delete it 

96 from the dictionary. 

97 """ 

98 with cls.__peers_by_uid_lock: 

99 cls.__peers_by_uid[peer_uid].__instance_count -= 1 

100 

101 if cls.__peers_by_uid[peer_uid].__instance_count < 1: 

102 del cls.__peers_by_uid[peer_uid]