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