diff --git a/src/msg/extensions/tools.py b/src/msg/extensions/tools.py index 8ce39f34..06db2f4e 100644 --- a/src/msg/extensions/tools.py +++ b/src/msg/extensions/tools.py @@ -19,6 +19,7 @@ 'headers':{'type':'object','maxProperties':30,'additionalProperties':{'type':'string','maxLength':4096}}, 'body':BYTES,'body_ref':REF},('url',)) TOOL_OUTPUT={'type':'object'} +MAX_TOOL_CONCURRENCY=32 def descriptor(name): @@ -49,11 +50,24 @@ async def resolve_tool_for_invoke(tx,value): return await resolve(tx,value) +def declared_concurrency(value): + # Published revisions predate this field; absence is the tightest limit. + limit=value.get('concurrency',1) + require(type(limit) is int and 1<=limit<=MAX_TOOL_CONCURRENCY,'invalid_tool_concurrency') + return limit + + +async def tool_concurrency(app,tx,ref): + rev=await tx.revision(ref) + return declared_concurrency(loads(await app.contents.read_bytes(rev.content,limit=65536))) + + async def read_tool(app,tx,rid,revision=None): resource=await tx.resource(rid) require(resource.type=='tool' and resource.parent==TOOLS_SPACE,'not_a_tool') rev=await tx.revision(ResourceRef(id=rid,revision=revision)) value=loads(await app.contents.read_bytes(rev.content,limit=65536)) + declared_concurrency(value) registered=app.registry.tool(rid,value['version']) require(value['tool_id']==rid and value['executor_key']==registered.executor_key, 'untrusted_tool_executor') diff --git a/src/msg/workers/effects.py b/src/msg/workers/effects.py index 3e48cd73..78e943b2 100644 --- a/src/msg/workers/effects.py +++ b/src/msg/workers/effects.py @@ -109,26 +109,45 @@ def __init__(self, app, *, tool_runner: ToolRunner | None = None, mail_sender=No self.lease_seconds = lease_seconds async def _claim(self): + from msg.extensions.tools import tool_concurrency async with self.app.metadata.transaction(write=True) as tx: self.app.runtime_generation.require_current(tx) if quarantine_active(tx): return None, False # Orphaned external work cannot be assumed not to have executed. + running = {} for (raw_id,) in tx.execute("SELECT id FROM jobs WHERE state='running' ORDER BY next_at,id"): job = await tx.job(raw_id) if job.lease_until is None or job.lease_until <= self.app.clock(): await tx.save_job(replace(job, state='uncertain', lease_until=None)) tx.set_setting('job_status:' + job.id, {'code': 'expired_execution_lease'}) return job, False - row = tx.one("SELECT id FROM jobs WHERE state='pending' AND next_at<=? ORDER BY next_at,id LIMIT 1", - (wire(self.app.clock()),)) - if row is None: - return None, False - job = await tx.job(row[0]) - job = replace(job, state='running', attempts=job.attempts+1, - lease_until=self.app.clock()+timedelta(seconds=self.lease_seconds)) - await tx.save_job(job) - return job, True + if job.kind == 'tool': + running[job.arguments['tool']['id']] = running.get(job.arguments['tool']['id'], 0) + 1 + # Other workers may hold this tool's jobs; only the running rows seen + # under this write transaction are a shared fact, never a local lock. + limits = {} + for (raw_id,) in tx.execute("SELECT id FROM jobs WHERE state='pending' AND next_at<=? ORDER BY next_at,id", + (wire(self.app.clock()),)): + job = await tx.job(raw_id) + if job.kind == 'tool': + ref = decode(ResourceRef, job.arguments['tool']) + if (ref.id, ref.revision) not in limits: + try: + limits[ref.id, ref.revision] = await tool_concurrency(self.app, tx, ref) + except Failure as exc: + if exc.retryable: + raise + await tx.save_job(replace(job, state='failed', lease_until=None)) + tx.set_setting('job_status:' + job.id, {'code': exc.code}) + return job, False + if running.get(ref.id, 0) >= limits[ref.id, ref.revision]: + continue + job = replace(job, state='running', attempts=job.attempts+1, + lease_until=self.app.clock()+timedelta(seconds=self.lease_seconds)) + await tx.save_job(job) + return job, True + return None, False async def _finish(self, job, state, code, *, status=None): async with self.app.metadata.transaction(write=True) as tx: diff --git a/tests/test_tools.py b/tests/test_tools.py index 2a3a6100..25ea8ec7 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -4,9 +4,11 @@ import pytest import httpx -from msg.core.codec import wire,canonical,digest +from msg.application import Application +from msg.core.codec import wire,canonical,digest,loads from msg.core.errors import Failure -from msg.core.models import Scope,NetworkPolicy,ResourceRef +from msg.core.models import Scope,NetworkPolicy,ResourceRef,Revision +from msg.plugins.common import new_id from msg.security.network import validate_url,validate_addresses,intersect_policy from msg.workers.effects import EffectWorker,ToolResult from msg.extensions.tools import descriptor, read_tool @@ -121,6 +123,152 @@ async def runner(tool,args,policies,outdir): assert (await tx.job(result.data['job_id'])).state=='uncertain' +async def _claimed_dns_job(app,root,handle): + key,uid,_=await register(app,handle) + cert=await approve(app,root,uid,key,(scoped(app,'tool.use','tool_dns',app.registry.capability('tool.use').operations),)) + result=await call(app,'tool.run',{'id':'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/tools/dns','arguments':{'name':'example.org','type':'A'}}, + key=key,subject=uid,certs=(cert.resource_id,)) + assert result.status=='accepted',wire(result) + async with app.metadata.transaction(write=False) as tx: + return await tx.job(result.data['job_id']) + + +async def _enqueue_like(app,template,seconds_ago,**changes): + job_id=new_id('job') + job=replace(template,id=job_id,dedupe_key='test:'+job_id,state='pending',attempts=0,lease_until=None, + next_attempt_at=NOW-timedelta(seconds=seconds_ago),**changes) + async with app.metadata.transaction(write=True) as tx: + await tx.enqueue(job) + return job + + +async def _tool_ref(app,rid): + async with app.metadata.transaction(write=False) as tx: + return wire(ResourceRef(id=rid,revision=(await tx.resource(rid)).revision)) + + +async def _dns_revision_declaring(app,concurrency): + async with app.metadata.transaction(write=True) as tx: + resource=await tx.resource('tool_dns') + blob=await app.contents.put_bytes(canonical({**descriptor('dns'),'concurrency':concurrency}),'application/json') + revision=Revision(format_version=1,id=new_id('v'),resource_id=resource.id,parents=(resource.revision,), + content=blob,relations=(),actor=resource.owner,subject=resource.owner,author=resource.owner, + created_at=app.clock(),manifest_digest='') + revision=replace(revision,manifest_digest=digest({k:v for k,v in wire(revision).items() + if k not in {'manifest_digest','signature'}})) + await app.contents.pin(blob,revision.id) + await app.contents.commit_revision(resource.parent,revision) + await tx.append_revision(revision) + return wire(ResourceRef(id=resource.id,revision=revision.id)) + + +async def _state(app,job_id): + async with app.metadata.transaction(write=False) as tx: + return (await tx.job(job_id)).state + + +@pytest.mark.asyncio +async def test_historical_tool_revision_bytes_have_no_concurrency_and_mean_one(installed): + app,_=installed + assert 'concurrency' not in descriptor('dns') + async with app.metadata.transaction(write=False) as tx: + rev=await tx.revision(ResourceRef(id='tool_dns')) + stored=await app.contents.read_bytes(rev.content,limit=65536) + assert stored==canonical(descriptor('dns')) + assert 'concurrency' not in loads(stored) + + +@pytest.mark.asyncio +async def test_same_tool_is_not_claimed_by_a_second_worker_while_other_jobs_proceed(installed): + app,root=installed + held=await _claimed_dns_job(app,root,'concurrency-holder') + first=EffectWorker(app,tool_runner=lambda *a:None) + claimed,execute=await first._claim() + assert execute and claimed.id==held.id + same_tool=await _enqueue_like(app,held,3) + other_tool=await _enqueue_like(app,held,2,arguments={**held.arguments,'tool':await _tool_ref(app,'tool_curl')}) + non_tool=await _enqueue_like(app,held,1,kind='mail',operation='communication.send', + arguments={'recipient_subject':held.principal.subject}) + peer=await Application(app.settings,clock=lambda:NOW).load() + try: + second=EffectWorker(peer,tool_runner=lambda *a:None) + order=[] + for _ in range(2): + job,execute=await second._claim() + assert execute + order.append(job.id) + assert order==[other_tool.id,non_tool.id] + assert await second._claim()==(None,False) + assert await _state(app,same_tool.id)=='pending' + await first._finish(claimed,'done','ok') + job,execute=await second._claim() + assert execute and job.id==same_tool.id + finally: + await peer.close() + + +@pytest.mark.asyncio +async def test_explicit_tool_concurrency_admits_up_to_the_declared_limit(installed): + app,root=installed + template=await _claimed_dns_job(app,root,'concurrency-two') + ref=await _dns_revision_declaring(app,2) + async with app.metadata.transaction(write=False) as tx: + assert (await read_tool(app,tx,'tool_dns',ref['revision'])).resource.revision==ref['revision'] + async with app.metadata.transaction(write=True) as tx: + await tx.save_job(replace(template,state='done')) + jobs=[await _enqueue_like(app,template,3-i,arguments={**template.arguments,'tool':ref}) for i in range(3)] + peer=await Application(app.settings,clock=lambda:NOW).load() + try: + workers=(EffectWorker(app,tool_runner=lambda *a:None),EffectWorker(peer,tool_runner=lambda *a:None)) + for worker,job in zip(workers,jobs): + claimed,execute=await worker._claim() + assert execute and claimed.id==job.id + for worker in workers: + assert await worker._claim()==(None,False) + assert await _state(app,jobs[2].id)=='pending' + finally: + await peer.close() + + +@pytest.mark.asyncio +async def test_historical_limit_counts_running_jobs_across_revisions_of_one_tool(installed): + app,root=installed + template=await _claimed_dns_job(app,root,'concurrency-mixed') + ref=await _dns_revision_declaring(app,2) + worker=EffectWorker(app,tool_runner=lambda *a:None) + claimed,_=await worker._claim() + assert claimed.id==template.id + historical=await _enqueue_like(app,template,2) + declared=await _enqueue_like(app,template,1,arguments={**template.arguments,'tool':ref}) + job,execute=await worker._claim() + assert execute and job.id==declared.id + assert await worker._claim()==(None,False) + assert await _state(app,historical.id)=='pending' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('declared',[0,33,-1,True,'2',1.5,None]) +async def test_invalid_declared_tool_concurrency_is_rejected_not_unbounded(installed,declared): + app,root=installed + template=await _claimed_dns_job(app,root,'concurrency-invalid') + ref=await _dns_revision_declaring(app,declared) + async with app.metadata.transaction(write=False) as tx: + with pytest.raises(Failure,match='invalid_tool_concurrency'): + await read_tool(app,tx,'tool_dns',ref['revision']) + async with app.metadata.transaction(write=True) as tx: + await tx.save_job(replace(template,state='done')) + bad=await _enqueue_like(app,template,1,arguments={**template.arguments,'tool':ref}) + ran=[] + async def runner(*args):ran.append(args) + worker=EffectWorker(app,tool_runner=runner) + job,execute=await worker._claim() + assert job.id==bad.id and not execute + assert await _state(app,bad.id)=='failed' and not ran + async with app.metadata.transaction(write=False) as tx: + assert tx.setting('job_status:'+bad.id)=={'code':'invalid_tool_concurrency'} + assert await worker._claim()==(None,False) + + @pytest.mark.asyncio async def test_tool_run_is_the_only_public_operation_and_old_short_code_is_retired(installed): app, _ = installed