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