##// END OF EJS Templates
path: pass `path` to `peer` in `hg fastannotate`...
marmoute -
r50630:8c6f7839 default
parent child Browse files
Show More
@@ -1,259 +1,262 b''
1 1 # Copyright 2016-present Facebook. All Rights Reserved.
2 2 #
3 3 # protocol: logic for a server providing fastannotate support
4 4 #
5 5 # This software may be used and distributed according to the terms of the
6 6 # GNU General Public License version 2 or any later version.
7 7
8 8 import contextlib
9 9 import os
10 10
11 11 from mercurial.i18n import _
12 12 from mercurial.pycompat import open
13 13 from mercurial import (
14 14 error,
15 15 extensions,
16 16 hg,
17 17 util,
18 18 wireprotov1peer,
19 19 wireprotov1server,
20 20 )
21 21 from mercurial.utils import (
22 22 urlutil,
23 23 )
24 24 from . import context
25 25
26 26 # common
27 27
28 28
29 29 def _getmaster(ui):
30 30 """get the mainbranch, and enforce it is set"""
31 31 master = ui.config(b'fastannotate', b'mainbranch')
32 32 if not master:
33 33 raise error.Abort(
34 34 _(
35 35 b'fastannotate.mainbranch is required '
36 36 b'for both the client and the server'
37 37 )
38 38 )
39 39 return master
40 40
41 41
42 42 # server-side
43 43
44 44
45 45 def _capabilities(orig, repo, proto):
46 46 result = orig(repo, proto)
47 47 result.append(b'getannotate')
48 48 return result
49 49
50 50
51 51 def _getannotate(repo, proto, path, lastnode):
52 52 # output:
53 53 # FILE := vfspath + '\0' + str(size) + '\0' + content
54 54 # OUTPUT := '' | FILE + OUTPUT
55 55 result = b''
56 56 buildondemand = repo.ui.configbool(
57 57 b'fastannotate', b'serverbuildondemand', True
58 58 )
59 59 with context.annotatecontext(repo, path) as actx:
60 60 if buildondemand:
61 61 # update before responding to the client
62 62 master = _getmaster(repo.ui)
63 63 try:
64 64 if not actx.isuptodate(master):
65 65 actx.annotate(master, master)
66 66 except Exception:
67 67 # non-fast-forward move or corrupted. rebuild automically.
68 68 actx.rebuild()
69 69 try:
70 70 actx.annotate(master, master)
71 71 except Exception:
72 72 actx.rebuild() # delete files
73 73 finally:
74 74 # although the "with" context will also do a close/flush, we
75 75 # need to do it early so we can send the correct respond to
76 76 # client.
77 77 actx.close()
78 78 # send back the full content of revmap and linelog, in the future we
79 79 # may want to do some rsync-like fancy updating.
80 80 # the lastnode check is not necessary if the client and the server
81 81 # agree where the main branch is.
82 82 if actx.lastnode != lastnode:
83 83 for p in [actx.revmappath, actx.linelogpath]:
84 84 if not os.path.exists(p):
85 85 continue
86 86 with open(p, b'rb') as f:
87 87 content = f.read()
88 88 vfsbaselen = len(repo.vfs.base + b'/')
89 89 relpath = p[vfsbaselen:]
90 90 result += b'%s\0%d\0%s' % (relpath, len(content), content)
91 91 return result
92 92
93 93
94 94 def _registerwireprotocommand():
95 95 if b'getannotate' in wireprotov1server.commands:
96 96 return
97 97 wireprotov1server.wireprotocommand(b'getannotate', b'path lastnode')(
98 98 _getannotate
99 99 )
100 100
101 101
102 102 def serveruisetup(ui):
103 103 _registerwireprotocommand()
104 104 extensions.wrapfunction(wireprotov1server, b'_capabilities', _capabilities)
105 105
106 106
107 107 # client-side
108 108
109 109
110 110 def _parseresponse(payload):
111 111 result = {}
112 112 i = 0
113 113 l = len(payload) - 1
114 114 state = 0 # 0: vfspath, 1: size
115 115 vfspath = size = b''
116 116 while i < l:
117 117 ch = payload[i : i + 1]
118 118 if ch == b'\0':
119 119 if state == 1:
120 120 result[vfspath] = payload[i + 1 : i + 1 + int(size)]
121 121 i += int(size)
122 122 state = 0
123 123 vfspath = size = b''
124 124 elif state == 0:
125 125 state = 1
126 126 else:
127 127 if state == 1:
128 128 size += ch
129 129 elif state == 0:
130 130 vfspath += ch
131 131 i += 1
132 132 return result
133 133
134 134
135 135 def peersetup(ui, peer):
136 136 class fastannotatepeer(peer.__class__):
137 137 @wireprotov1peer.batchable
138 138 def getannotate(self, path, lastnode=None):
139 139 if not self.capable(b'getannotate'):
140 140 ui.warn(_(b'remote peer cannot provide annotate cache\n'))
141 141 return None, None
142 142 else:
143 143 args = {b'path': path, b'lastnode': lastnode or b''}
144 144 return args, _parseresponse
145 145
146 146 peer.__class__ = fastannotatepeer
147 147
148 148
149 149 @contextlib.contextmanager
150 150 def annotatepeer(repo):
151 151 ui = repo.ui
152 152
153 153 remotedest = ui.config(b'fastannotate', b'remotepath', b'default')
154 r = urlutil.get_unique_pull_path(b'fastannotate', repo, ui, remotedest)
155 remotepath = r[0]
154 remotepath = urlutil.get_unique_pull_path_obj(
155 b'fastannotate',
156 ui,
157 remotedest,
158 )
156 159 peer = hg.peer(ui, {}, remotepath)
157 160
158 161 try:
159 162 yield peer
160 163 finally:
161 164 peer.close()
162 165
163 166
164 167 def clientfetch(repo, paths, lastnodemap=None, peer=None):
165 168 """download annotate cache from the server for paths"""
166 169 if not paths:
167 170 return
168 171
169 172 if peer is None:
170 173 with annotatepeer(repo) as peer:
171 174 return clientfetch(repo, paths, lastnodemap, peer)
172 175
173 176 if lastnodemap is None:
174 177 lastnodemap = {}
175 178
176 179 ui = repo.ui
177 180 results = []
178 181 with peer.commandexecutor() as batcher:
179 182 ui.debug(b'fastannotate: requesting %d files\n' % len(paths))
180 183 for p in paths:
181 184 results.append(
182 185 batcher.callcommand(
183 186 b'getannotate',
184 187 {b'path': p, b'lastnode': lastnodemap.get(p)},
185 188 )
186 189 )
187 190
188 191 for result in results:
189 192 r = result.result()
190 193 # TODO: pconvert these paths on the server?
191 194 r = {util.pconvert(p): v for p, v in r.items()}
192 195 for path in sorted(r):
193 196 # ignore malicious paths
194 197 if not path.startswith(b'fastannotate/') or b'/../' in (
195 198 path + b'/'
196 199 ):
197 200 ui.debug(
198 201 b'fastannotate: ignored malicious path %s\n' % path
199 202 )
200 203 continue
201 204 content = r[path]
202 205 if ui.debugflag:
203 206 ui.debug(
204 207 b'fastannotate: writing %d bytes to %s\n'
205 208 % (len(content), path)
206 209 )
207 210 repo.vfs.makedirs(os.path.dirname(path))
208 211 with repo.vfs(path, b'wb') as f:
209 212 f.write(content)
210 213
211 214
212 215 def _filterfetchpaths(repo, paths):
213 216 """return a subset of paths whose history is long and need to fetch linelog
214 217 from the server. works with remotefilelog and non-remotefilelog repos.
215 218 """
216 219 threshold = repo.ui.configint(b'fastannotate', b'clientfetchthreshold', 10)
217 220 if threshold <= 0:
218 221 return paths
219 222
220 223 result = []
221 224 for path in paths:
222 225 try:
223 226 if len(repo.file(path)) >= threshold:
224 227 result.append(path)
225 228 except Exception: # file not found etc.
226 229 result.append(path)
227 230
228 231 return result
229 232
230 233
231 234 def localreposetup(ui, repo):
232 235 class fastannotaterepo(repo.__class__):
233 236 def prefetchfastannotate(self, paths, peer=None):
234 237 master = _getmaster(self.ui)
235 238 needupdatepaths = []
236 239 lastnodemap = {}
237 240 try:
238 241 for path in _filterfetchpaths(self, paths):
239 242 with context.annotatecontext(self, path) as actx:
240 243 if not actx.isuptodate(master, strict=False):
241 244 needupdatepaths.append(path)
242 245 lastnodemap[path] = actx.lastnode
243 246 if needupdatepaths:
244 247 clientfetch(self, needupdatepaths, lastnodemap, peer)
245 248 except Exception as ex:
246 249 # could be directory not writable or so, not fatal
247 250 self.ui.debug(b'fastannotate: prefetch failed: %r\n' % ex)
248 251
249 252 repo.__class__ = fastannotaterepo
250 253
251 254
252 255 def clientreposetup(ui, repo):
253 256 _registerwireprotocommand()
254 257 if repo.local():
255 258 localreposetup(ui, repo)
256 259 # TODO: this mutates global state, but only if at least one repo
257 260 # has the extension enabled. This is probably bad for hgweb.
258 261 if peersetup not in hg.wirepeersetupfuncs:
259 262 hg.wirepeersetupfuncs.append(peersetup)
General Comments 0
You need to be logged in to leave comments. Login now