Differential D6984 Diff 25406 swh/loader/mercurial/hgutil.py

Changeset View

Standalone View

swh/loader/mercurial/hgutil.py

# Copyright (C) 2020-2021 The Software Heritage developers		# Copyright (C) 2020-2021 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution		# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version		# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information		# See top-level LICENSE file for more information

from collections import defaultdict		from collections import defaultdict
from dataclasses import dataclass		from dataclasses import dataclass
import io		from functools import partial
import os
import signal
import time
import traceback
from typing import Dict, List, Mapping, NewType, Optional, Set		from typing import Dict, List, Mapping, NewType, Optional, Set

from billiard import Process, Queue

# The internal Mercurial API is not guaranteed to be stable.		# The internal Mercurial API is not guaranteed to be stable.
from mercurial import bookmarks, context, error, hg, smartset, util # type: ignore		from mercurial import bookmarks, context, error, hg, smartset, util # type: ignore
import mercurial.ui # type: ignore		import mercurial.ui # type: ignore

		from swh.loader.core.utils import clone_with_timeout

NULLID = mercurial.node.nullid		NULLID = mercurial.node.nullid
HgNodeId = NewType("HgNodeId", bytes)		HgNodeId = NewType("HgNodeId", bytes)
Repository = hg.localrepo		Repository = hg.localrepo
BaseContext = context.basectx		BaseContext = context.basectx
LRUCacheDict = util.lrucachedict		LRUCacheDict = util.lrucachedict
HgSpanSet = smartset._spanset		HgSpanSet = smartset._spanset
HgFilteredSet = smartset.filteredset		HgFilteredSet = smartset.filteredset
LookupError = error.LookupError		LookupError = error.LookupError
▲ Show 20 Lines • Show All 82 Lines • ▼ Show 20 Lines	return BranchingInfo(
tips=branch_tips,		tips=branch_tips,
bookmarks=branch_bookmarks,		bookmarks=branch_bookmarks,
open_heads=branch_open_heads,		open_heads=branch_open_heads,
closed_heads=branch_closed_heads,		closed_heads=branch_closed_heads,
default_branch_alias=default_rev_alias,		default_branch_alias=default_rev_alias,
)		)


class CloneTimeout(Exception):		def clone(src: str, dest: str, timeout: float):
pass		closure = partial(
		hg.clone,

class CloneFailure(Exception):
pass


def _clone_task(src: str, dest: str, errors: Queue) -> None:
"""Clone task to run in a subprocess.

Args:
src: clone source
dest: clone destination
errors: message queue to communicate errors
"""
try:
hg.clone(
ui=mercurial.ui.ui.load(),		ui=mercurial.ui.ui.load(),
peeropts={},		peeropts={},
source=src.encode(),		source=src.encode(),
dest=dest.encode(),		dest=dest.encode(),
update=False,		update=False,
)		)
except Exception as e:		clone_with_timeout(src, dest, closure, timeout)
exc_buffer = io.StringIO()
traceback.print_exc(file=exc_buffer)
errors.put_nowait(exc_buffer.getvalue())
raise e


def clone(src: str, dest: str, timeout: float) -> None:
"""Clone a repository with timeout.

Args:
src: clone source
dest: clone destination
timeout: timeout in seconds
"""
errors: Queue = Queue()
process = Process(target=_clone_task, args=(src, dest, errors))
process.start()
process.join(timeout)

if process.is_alive():
process.terminate()
# Give it literally a second (in successive steps of 0.1 second), then kill it.
# Can't use `process.join(1)` here, billiard appears to be bugged
# https://github.com/celery/billiard/issues/270
killed = False
for _ in range(10):
time.sleep(0.1)
if not process.is_alive():
break
else:
killed = True
os.kill(process.pid, signal.SIGKILL)
raise CloneTimeout(src, timeout, killed)

if not errors.empty():
raise CloneFailure(src, dest, errors.get())