waf

wafcache.py (18813B)

---

 #! /usr/bin/env python
 # encoding: utf-8
 # Thomas Nagy, 2019 (ita)
 
 """
 Filesystem-based cache system to share and re-use build artifacts
 
 Cache access operations (copy to and from) are delegated to
 independent pre-forked worker subprocesses.
 
 The following environment variables may be set:
 * WAFCACHE: several possibilities:
   - File cache:
     absolute path of the waf cache (~/.cache/wafcache_user,
     where `user` represents the currently logged-in user)
   - URL to a cache server, for example:
     export WAFCACHE=http://localhost:8080/files/
     in that case, GET/POST requests are made to urls of the form
     http://localhost:8080/files/000000000/0 (cache management is delegated to the server)
   - GCS, S3 or MINIO bucket
     gs://my-bucket/    (uses gsutil command line tool or WAFCACHE_CMD)
     s3://my-bucket/    (uses aws command line tool or WAFCACHE_CMD)
     minio://my-bucket/ (uses mc command line tool or WAFCACHE_CMD)
 * WAFCACHE_CMD: bucket upload/download command, for example:
     WAFCACHE_CMD="gsutil cp %{SRC} %{TGT}"
   Note that the WAFCACHE bucket value is used for the source or destination
   depending on the operation (upload or download). For example, with:
     WAFCACHE="gs://mybucket/"
   the following commands may be run:
     gsutil cp build/myprogram  gs://mybucket/aa/aaaaa/1
     gsutil cp gs://mybucket/bb/bbbbb/2 build/somefile
 * WAFCACHE_NO_PUSH: if set, disables pushing to the cache
 * WAFCACHE_VERBOSITY: if set, displays more detailed cache operations
 * WAFCACHE_STATS: if set, displays cache usage statistics on exit
 
 File cache specific options:
   Files are copied using hard links by default; if the cache is located
   onto another partition, the system switches to file copies instead.
 * WAFCACHE_TRIM_MAX_FOLDER: maximum amount of tasks to cache (1M)
 * WAFCACHE_EVICT_MAX_BYTES: maximum amount of cache size in bytes (10GB)
 * WAFCACHE_EVICT_INTERVAL_MINUTES: minimum time interval to try
                                    and trim the cache (3 minutes)
 
 Upload specific options:
 * WAFCACHE_ASYNC_WORKERS: define a number of workers to upload results asynchronously
                           this may improve build performance with many/long file uploads
                           the default is unset (synchronous uploads)
 * WAFCACHE_ASYNC_NOWAIT: do not wait for uploads to complete (default: False)
                          this requires asynchonous uploads to have an effect
 
 Usage::
 
 	def build(bld):
 		bld.load('wafcache')
 		...
 
 To troubleshoot::
 
 	waf clean build --zone=wafcache
 """
 
 import atexit, base64, errno, fcntl, getpass, os, re, shutil, sys, time, threading, traceback, urllib3, shlex
 try:
 	import subprocess32 as subprocess
 except ImportError:
 	import subprocess
 
 base_cache = os.path.expanduser('~/.cache/')
 if not os.path.isdir(base_cache):
 	base_cache = '/tmp/'
 default_wafcache_dir = os.path.join(base_cache, 'wafcache_' + getpass.getuser())
 
 CACHE_DIR = os.environ.get('WAFCACHE', default_wafcache_dir)
 WAFCACHE_CMD = os.environ.get('WAFCACHE_CMD')
 TRIM_MAX_FOLDERS = int(os.environ.get('WAFCACHE_TRIM_MAX_FOLDER', 1000000))
 EVICT_INTERVAL_MINUTES = int(os.environ.get('WAFCACHE_EVICT_INTERVAL_MINUTES', 3))
 EVICT_MAX_BYTES = int(os.environ.get('WAFCACHE_EVICT_MAX_BYTES', 10**10))
 WAFCACHE_NO_PUSH = 1 if os.environ.get('WAFCACHE_NO_PUSH') else 0
 WAFCACHE_VERBOSITY = 1 if os.environ.get('WAFCACHE_VERBOSITY') else 0
 WAFCACHE_STATS = 1 if os.environ.get('WAFCACHE_STATS') else 0
 WAFCACHE_ASYNC_WORKERS = os.environ.get('WAFCACHE_ASYNC_WORKERS')
 WAFCACHE_ASYNC_NOWAIT = os.environ.get('WAFCACHE_ASYNC_NOWAIT')
 OK = "ok"
 
 re_waf_cmd = re.compile('(?P<src>%{SRC})|(?P<tgt>%{TGT})')
 
 try:
 	import cPickle
 except ImportError:
 	import pickle as cPickle
 
 if __name__ != '__main__':
 	from waflib import Task, Logs, Utils, Build
 
 def can_retrieve_cache(self):
 	"""
 	New method for waf Task classes
 	"""
 	if not self.outputs:
 		return False
 
 	self.cached = False
 
 	sig = self.signature()
 	ssig = Utils.to_hex(self.uid() + sig)
 
 	if WAFCACHE_STATS:
 		self.generator.bld.cache_reqs += 1
 
 	files_to = [node.abspath() for node in self.outputs]
 	proc = get_process()
 	err = cache_command(proc, ssig, [], files_to)
 	process_pool.append(proc)
 	if err.startswith(OK):
 		if WAFCACHE_VERBOSITY:
 			Logs.pprint('CYAN', '  Fetched %r from cache' % files_to)
 		else:
 			Logs.debug('wafcache: fetched %r from cache', files_to)
 		if WAFCACHE_STATS:
 			self.generator.bld.cache_hits += 1
 	else:
 		if WAFCACHE_VERBOSITY:
 			Logs.pprint('YELLOW', '  No cache entry %s' % files_to)
 		else:
 			Logs.debug('wafcache: No cache entry %s: %s', files_to, err)
 		return False
 
 	self.cached = True
 	return True
 
 def put_files_cache(self):
 	"""
 	New method for waf Task classes
 	"""
 	if WAFCACHE_NO_PUSH or getattr(self, 'cached', None) or not self.outputs:
 		return
 
 	files_from = []
 	for node in self.outputs:
 		path = node.abspath()
 		if not os.path.isfile(path):
 			return
 		files_from.append(path)
 
 	bld = self.generator.bld
 	old_sig = self.signature()
 
 	for node in self.inputs:
 		try:
 			del node.ctx.cache_sig[node]
 		except KeyError:
 			pass
 
 	delattr(self, 'cache_sig')
 	sig = self.signature()
 
 	def _async_put_files_cache(bld, ssig, files_from):
 		proc = get_process()
 		if WAFCACHE_ASYNC_WORKERS:
 			with bld.wafcache_lock:
 				if bld.wafcache_stop:
 					process_pool.append(proc)
 					return
 				bld.wafcache_procs.add(proc)
 
 		err = cache_command(proc, ssig, files_from, [])
 		process_pool.append(proc)
 		if err.startswith(OK):
 			if WAFCACHE_VERBOSITY:
 				Logs.pprint('CYAN', '  Successfully uploaded %s to cache' % files_from)
 			else:
 				Logs.debug('wafcache: Successfully uploaded %r to cache', files_from)
 			if WAFCACHE_STATS:
 				bld.cache_puts += 1
 		else:
 			if WAFCACHE_VERBOSITY:
 				Logs.pprint('RED', '  Error caching step results %s: %s' % (files_from, err))
 			else:
 				Logs.debug('wafcache: Error caching results %s: %s', files_from, err)
 
 	if old_sig == sig:
 		ssig = Utils.to_hex(self.uid() + sig)
 		if WAFCACHE_ASYNC_WORKERS:
 			fut = bld.wafcache_executor.submit(_async_put_files_cache, bld, ssig, files_from)
 			bld.wafcache_uploads.append(fut)
 		else:
 			_async_put_files_cache(bld, ssig, files_from)
 	else:
 		Logs.debug('wafcache: skipped %r upload due to late input modifications %r', self.outputs, self.inputs)
 
 	bld.task_sigs[self.uid()] = self.cache_sig
 
 def hash_env_vars(self, env, vars_lst):
 	"""
 	Reimplement BuildContext.hash_env_vars so that the resulting hash does not depend on local paths
 	"""
 	if not env.table:
 		env = env.parent
 		if not env:
 			return Utils.SIG_NIL
 
 	idx = str(id(env)) + str(vars_lst)
 	try:
 		cache = self.cache_env
 	except AttributeError:
 		cache = self.cache_env = {}
 	else:
 		try:
 			return self.cache_env[idx]
 		except KeyError:
 			pass
 
 	v = str([env[a] for a in vars_lst])
 	v = v.replace(self.srcnode.abspath().__repr__()[:-1], '')
 	m = Utils.md5()
 	m.update(v.encode())
 	ret = m.digest()
 
 	Logs.debug('envhash: %r %r', ret, v)
 
 	cache[idx] = ret
 
 	return ret
 
 def uid(self):
 	"""
 	Reimplement Task.uid() so that the signature does not depend on local paths
 	"""
 	try:
 		return self.uid_
 	except AttributeError:
 		m = Utils.md5()
 		src = self.generator.bld.srcnode
 		up = m.update
 		up(self.__class__.__name__.encode())
 		for x in self.inputs + self.outputs:
 			up(x.path_from(src).encode())
 		self.uid_ = m.digest()
 		return self.uid_
 
 
 def make_cached(cls):
 	"""
 	Enable the waf cache for a given task class
 	"""
 	if getattr(cls, 'nocache', None) or getattr(cls, 'has_cache', False):
 		return
 
 	full_name = "%s.%s" % (cls.__module__, cls.__name__)
 	if full_name in ('waflib.Tools.ccroot.vnum', 'waflib.Build.inst'):
 		return
 
 	m1 = getattr(cls, 'run', None)
 	def run(self):
 		if getattr(self, 'nocache', False):
 			return m1(self)
 		if self.can_retrieve_cache():
 			return 0
 		return m1(self)
 	cls.run = run
 
 	m2 = getattr(cls, 'post_run', None)
 	def post_run(self):
 		if getattr(self, 'nocache', False):
 			return m2(self)
 		ret = m2(self)
 		self.put_files_cache()
 		return ret
 	cls.post_run = post_run
 	cls.has_cache = True
 
 process_pool = []
 def get_process():
 	"""
 	Returns a worker process that can process waf cache commands
 	The worker process is assumed to be returned to the process pool when unused
 	"""
 	try:
 		return process_pool.pop()
 	except IndexError:
 		filepath = os.path.dirname(os.path.abspath(__file__)) + os.sep + 'wafcache.py'
 		cmd = [sys.executable, '-c', Utils.readf(filepath)]
 		return subprocess.Popen(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
 
 def atexit_pool():
 	for proc in process_pool:
 		proc.kill()
 atexit.register(atexit_pool)
 
 def build(bld):
 	"""
 	Called during the build process to enable file caching
 	"""
 
 	if WAFCACHE_ASYNC_WORKERS:
 		try:
 			num_workers = int(WAFCACHE_ASYNC_WORKERS)
 		except ValueError:
 			Logs.warn('Invalid WAFCACHE_ASYNC_WORKERS specified: %r' % WAFCACHE_ASYNC_WORKERS)
 		else:
 			from concurrent.futures import ThreadPoolExecutor
 			bld.wafcache_executor = ThreadPoolExecutor(max_workers=num_workers)
 			bld.wafcache_uploads = []
 			bld.wafcache_procs = set([])
 			bld.wafcache_stop = False
 			bld.wafcache_lock = threading.Lock()
 
 		def finalize_upload_async(bld):
 			if WAFCACHE_ASYNC_NOWAIT:
 				with bld.wafcache_lock:
 					bld.wafcache_stop = True
 
 				for fut in reversed(bld.wafcache_uploads):
 					fut.cancel()
 
 				for proc in bld.wafcache_procs:
 					proc.kill()
 
 				bld.wafcache_procs.clear()
 			else:
 				Logs.pprint('CYAN', '... waiting for wafcache uploads to complete (%s uploads)' % len(bld.wafcache_uploads))
 			bld.wafcache_executor.shutdown(wait=True)
 		bld.add_post_fun(finalize_upload_async)
 
 	if WAFCACHE_STATS:
 		# Init counter for statistics and hook to print results at the end
 		bld.cache_reqs = bld.cache_hits = bld.cache_puts = 0
 
 		def printstats(bld):
 			hit_ratio = 0
 			if bld.cache_reqs > 0:
 				hit_ratio = (bld.cache_hits / bld.cache_reqs) * 100
 			Logs.pprint('CYAN', '  wafcache stats: %s requests, %s hits (ratio: %.2f%%), %s writes' %
 					 (bld.cache_reqs, bld.cache_hits, hit_ratio, bld.cache_puts) )
 		bld.add_post_fun(printstats)
 
 	if process_pool:
 		# already called once
 		return
 
 	# pre-allocation
 	processes = [get_process() for x in range(bld.jobs)]
 	process_pool.extend(processes)
 
 	Task.Task.can_retrieve_cache = can_retrieve_cache
 	Task.Task.put_files_cache = put_files_cache
 	Task.Task.uid = uid
 	Build.BuildContext.hash_env_vars = hash_env_vars
 	for x in reversed(list(Task.classes.values())):
 		make_cached(x)
 
 def cache_command(proc, sig, files_from, files_to):
 	"""
 	Create a command for cache worker processes, returns a pickled
 	base64-encoded tuple containing the task signature, a list of files to
 	cache and a list of files files to get from cache (one of the lists
 	is assumed to be empty)
 	"""
 	obj = base64.b64encode(cPickle.dumps([sig, files_from, files_to]))
 	proc.stdin.write(obj)
 	proc.stdin.write('\n'.encode())
 	proc.stdin.flush()
 	obj = proc.stdout.readline()
 	if not obj:
 		raise OSError('Preforked sub-process %r died' % proc.pid)
 	return cPickle.loads(base64.b64decode(obj))
 
 try:
 	copyfun = os.link
 except NameError:
 	copyfun = shutil.copy2
 
 def atomic_copy(orig, dest):
 	"""
 	Copy files to the cache, the operation is atomic for a given file
 	"""
 	global copyfun
 	tmp = dest + '.tmp'
 	up = os.path.dirname(dest)
 	try:
 		os.makedirs(up)
 	except OSError:
 		pass
 
 	try:
 		copyfun(orig, tmp)
 	except OSError as e:
 		if e.errno == errno.EXDEV:
 			copyfun = shutil.copy2
 			copyfun(orig, tmp)
 		else:
 			raise
 	os.rename(tmp, dest)
 
 def lru_trim():
 	"""
 	the cache folders take the form:
 	`CACHE_DIR/0b/0b180f82246d726ece37c8ccd0fb1cde2650d7bfcf122ec1f169079a3bfc0ab9`
 	they are listed in order of last access, and then removed
 	until the amount of folders is within TRIM_MAX_FOLDERS and the total space
 	taken by files is less than EVICT_MAX_BYTES
 	"""
 	lst = []
 	for up in os.listdir(CACHE_DIR):
 		if len(up) == 2:
 			sub = os.path.join(CACHE_DIR, up)
 			for hval in os.listdir(sub):
 				path = os.path.join(sub, hval)
 
 				size = 0
 				for fname in os.listdir(path):
 					try:
 						size += os.lstat(os.path.join(path, fname)).st_size
 					except OSError:
 						pass
 				lst.append((os.stat(path).st_mtime, size, path))
 
 	lst.sort(key=lambda x: x[0])
 	lst.reverse()
 
 	tot = sum(x[1] for x in lst)
 	while tot > EVICT_MAX_BYTES or len(lst) > TRIM_MAX_FOLDERS:
 		_, tmp_size, path = lst.pop()
 		tot -= tmp_size
 
 		tmp = path + '.remove'
 		try:
 			shutil.rmtree(tmp)
 		except OSError:
 			pass
 		try:
 			os.rename(path, tmp)
 		except OSError:
 			sys.stderr.write('Could not rename %r to %r\n' % (path, tmp))
 		else:
 			try:
 				shutil.rmtree(tmp)
 			except OSError:
 				sys.stderr.write('Could not remove %r\n' % tmp)
 	sys.stderr.write("Cache trimmed: %r bytes in %r folders left\n" % (tot, len(lst)))
 
 
 def lru_evict():
 	"""
 	Reduce the cache size
 	"""
 	lockfile = os.path.join(CACHE_DIR, 'all.lock')
 	try:
 		st = os.stat(lockfile)
 	except EnvironmentError as e:
 		if e.errno == errno.ENOENT:
 			with open(lockfile, 'w') as f:
 				f.write('')
 			return
 		else:
 			raise
 
 	if st.st_mtime < time.time() - EVICT_INTERVAL_MINUTES * 60:
 		# check every EVICT_INTERVAL_MINUTES minutes if the cache is too big
 		# OCLOEXEC is unnecessary because no processes are spawned
 		fd = os.open(lockfile, os.O_RDWR | os.O_CREAT, 0o755)
 		try:
 			try:
 				fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
 			except EnvironmentError:
 				if WAFCACHE_VERBOSITY:
 					sys.stderr.write('wafcache: another cleaning process is running\n')
 			else:
 				# now dow the actual cleanup
 				lru_trim()
 				os.utime(lockfile, None)
 		finally:
 			os.close(fd)
 
 class netcache(object):
 	def __init__(self):
 		self.http = urllib3.PoolManager()
 
 	def url_of(self, sig, i):
 		return "%s/%s/%s" % (CACHE_DIR, sig, i)
 
 	def upload(self, file_path, sig, i):
 		url = self.url_of(sig, i)
 		with open(file_path, 'rb') as f:
 			file_data = f.read()
 		r = self.http.request('POST', url, timeout=60,
 			fields={ 'file': ('%s/%s' % (sig, i), file_data), })
 		if r.status >= 400:
 			raise OSError("Invalid status %r %r" % (url, r.status))
 
 	def download(self, file_path, sig, i):
 		url = self.url_of(sig, i)
 		with self.http.request('GET', url, preload_content=False, timeout=60) as inf:
 			if inf.status >= 400:
 				raise OSError("Invalid status %r %r" % (url, inf.status))
 			with open(file_path, 'wb') as out:
 				shutil.copyfileobj(inf, out)
 
 	def copy_to_cache(self, sig, files_from, files_to):
 		try:
 			for i, x in enumerate(files_from):
 				if not os.path.islink(x):
 					self.upload(x, sig, i)
 		except Exception:
 			return traceback.format_exc()
 		return OK
 
 	def copy_from_cache(self, sig, files_from, files_to):
 		try:
 			for i, x in enumerate(files_to):
 				self.download(x, sig, i)
 		except Exception:
 			return traceback.format_exc()
 		return OK
 
 class fcache(object):
 	def __init__(self):
 		if not os.path.exists(CACHE_DIR):
 			try:
 				os.makedirs(CACHE_DIR)
 			except OSError:
 				pass
 		if not os.path.exists(CACHE_DIR):
 			raise ValueError('Could not initialize the cache directory')
 
 	def copy_to_cache(self, sig, files_from, files_to):
 		"""
 		Copy files to the cache, existing files are overwritten,
 		and the copy is atomic only for a given file, not for all files
 		that belong to a given task object
 		"""
 		try:
 			for i, x in enumerate(files_from):
 				dest = os.path.join(CACHE_DIR, sig[:2], sig, str(i))
 				atomic_copy(x, dest)
 		except Exception:
 			return traceback.format_exc()
 		else:
 			# attempt trimming if caching was successful:
 			# we may have things to trim!
 			try:
 				lru_evict()
 			except Exception:
 				return traceback.format_exc()
 		return OK
 
 	def copy_from_cache(self, sig, files_from, files_to):
 		"""
 		Copy files from the cache
 		"""
 		try:
 			for i, x in enumerate(files_to):
 				orig = os.path.join(CACHE_DIR, sig[:2], sig, str(i))
 				atomic_copy(orig, x)
 
 			# success! update the cache time
 			os.utime(os.path.join(CACHE_DIR, sig[:2], sig), None)
 		except Exception:
 			return traceback.format_exc()
 		return OK
 
 class bucket_cache(object):
 	def bucket_copy(self, source, target):
 		if WAFCACHE_CMD:
 			def replacer(match):
 				if match.group('src'):
 					return source
 				elif match.group('tgt'):
 					return target
 			cmd = [re_waf_cmd.sub(replacer, x) for x in shlex.split(WAFCACHE_CMD)]
 		elif CACHE_DIR.startswith('s3://'):
 			cmd = ['aws', 's3', 'cp', source, target]
 		elif CACHE_DIR.startswith('gs://'):
 			cmd = ['gsutil', 'cp', source, target]
 		else:
 			cmd = ['mc', 'cp', source, target]
 
 		proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 		out, err = proc.communicate()
 		if proc.returncode:
 			raise OSError('Error copy %r to %r using: %r (exit %r):\n  out:%s\n  err:%s' % (
 				source, target, cmd, proc.returncode, out.decode(errors='replace'), err.decode(errors='replace')))
 
 	def copy_to_cache(self, sig, files_from, files_to):
 		try:
 			for i, x in enumerate(files_from):
 				dest = os.path.join(CACHE_DIR, sig[:2], sig, str(i))
 				self.bucket_copy(x, dest)
 		except Exception:
 			return traceback.format_exc()
 		return OK
 
 	def copy_from_cache(self, sig, files_from, files_to):
 		try:
 			for i, x in enumerate(files_to):
 				orig = os.path.join(CACHE_DIR, sig[:2], sig, str(i))
 				self.bucket_copy(orig, x)
 		except EnvironmentError:
 			return traceback.format_exc()
 		return OK
 
 def loop(service):
 	"""
 	This function is run when this file is run as a standalone python script,
 	it assumes a parent process that will communicate the commands to it
 	as pickled-encoded tuples (one line per command)
 
 	The commands are to copy files to the cache or copy files from the
 	cache to a target destination
 	"""
 	# one operation is performed at a single time by a single process
 	# therefore stdin never has more than one line
 	txt = sys.stdin.readline().strip()
 	if not txt:
 		# parent process probably ended
 		sys.exit(1)
 	ret = OK
 
 	[sig, files_from, files_to] = cPickle.loads(base64.b64decode(txt))
 	if files_from:
 		# TODO return early when pushing files upstream
 		ret = service.copy_to_cache(sig, files_from, files_to)
 	elif files_to:
 		# the build process waits for workers to (possibly) obtain files from the cache
 		ret = service.copy_from_cache(sig, files_from, files_to)
 	else:
 		ret = "Invalid command"
 
 	obj = base64.b64encode(cPickle.dumps(ret))
 	sys.stdout.write(obj.decode())
 	sys.stdout.write('\n')
 	sys.stdout.flush()
 
 if __name__ == '__main__':
 	if CACHE_DIR.startswith('s3://') or CACHE_DIR.startswith('gs://') or CACHE_DIR.startswith('minio://'):
 		if CACHE_DIR.startswith('minio://'):
 			CACHE_DIR = CACHE_DIR[8:]   # minio doesn't need the protocol part, uses config aliases
 		service = bucket_cache()
 	elif CACHE_DIR.startswith('http'):
 		service = netcache()
 	else:
 		service = fcache()
 	while 1:
 		try:
 			loop(service)
 		except KeyboardInterrupt:
 			break