Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

# Copyright (C) 2019 Bloomberg LP 

# 

# Licensed under the Apache License, Version 2.0 (the "License"); 

# you may not use this file except in compliance with the License. 

# You may obtain a copy of the License at 

# 

# <http://www.apache.org/licenses/LICENSE-2.0> 

# 

# Unless required by applicable law or agreed to in writing, software 

# distributed under the License is distributed on an "AS IS" BASIS, 

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

# See the License for the specific language governing permissions and 

# limitations under the License. 

 

import logging 

from threading import Lock 

 

from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 

from buildgrid.settings import REQUEST_METADATA_HEADER_NAME 

 

 

class Peer: 

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

 

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

__peers_by_uid = {} 

__peers_by_uid_lock = Lock() 

 

@classmethod 

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

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

 

Args: 

uid (str): a unique identifier 

token (str): an authentication token (optional) 

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

being registered 

 

Returns: 

Peer: an existing or newly created Peer object 

""" 

 

request_metadata = cls._context_extract_request_metadata(context) 

 

if request_metadata and request_metadata.tool_details: 

tool_name = request_metadata.tool_details.tool_name 

tool_version = request_metadata.tool_details.tool_version 

else: 

tool_name = tool_version = None 

 

with cls.__peers_by_uid_lock: 

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

# instance: 

if uid in cls.__peers_by_uid: 

existing_peer = cls.__peers_by_uid[uid] 

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

'of Peer with uid %s ', uid) 

existing_peer.__instance_count += 1 

return existing_peer 

else: 

# Otherwise we add ourselves to the list of Peers: 

new_peer = Peer(uid=uid, token=token, tool_name=tool_name, 

tool_version=tool_version) 

cls.__peers_by_uid[uid] = new_peer 

return cls.__peers_by_uid[uid] 

 

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

"""Creates a new Peer object. 

 

Args: 

uid (str): a unique identifier 

token (str): an authentication token (optional) 

tool_name(str): reported tool name for the peer's request 

tool_version(str): reported tool version for the peer's request 

""" 

self._uid = uid # This uniquely identifies a client 

self._token = token 

 

self._tool_name = tool_name 

self._tool_version = tool_version 

 

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

self.__instance_count = 1 

 

@classmethod 

def find_peer(cls, uid): 

return cls.__peers_by_uid.get(uid, None) 

 

def __eq__(self, other): 

if not isinstance(other, Peer): 

return False 

 

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

self.request_metadata == other.request_metadata 

 

def __hash__(self): 

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

 

def __str__(self): 

return 'Peer: uid: {}, tool_details: {} - {}'.format(self._uid, 

self._tool_name, 

self._tool_version) 

 

@property 

def uid(self): 

return self._uid 

 

@property 

def token(self): 

return self._token 

 

# -- `RequestMetadata` optional values (attached to the Execute() call) -- 

@property 

def tool_name(self): 

return self._tool_name 

 

@property 

def tool_version(self): 

return self._tool_version 

 

@classmethod 

def deregister_peer(cls, peer_uid): 

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

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

from the dictionary. 

""" 

with cls.__peers_by_uid_lock: 

cls.__peers_by_uid[peer_uid].__instance_count -= 1 

 

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

del cls.__peers_by_uid[peer_uid] 

 

@classmethod 

def _context_extract_request_metadata(cls, context): 

"""Given a `grpc.ServicerContext` object, extract the RequestMetadata 

header values if they are present. If they were not provided, 

returns None. 

 

Args: 

context (grpc.ServicerContext): Context for a RPC call. 

 

Returns: 

A `RequestMetadata` proto if RequestMetadata values are present, 

otherwise None. 

""" 

invocation_metadata = context.invocation_metadata() 

request_metadata_entry = next((entry for entry in invocation_metadata 

if entry.key == REQUEST_METADATA_HEADER_NAME), 

None) 

if not request_metadata_entry: 

return None 

 

request_metadata = remote_execution_pb2.RequestMetadata() 

request_metadata.ParseFromString(request_metadata_entry.value) 

 

return request_metadata