From 176edd04107d80e838666504715d79f411badfff Mon Sep 17 00:00:00 2001 From: George Pollard Date: Tue, 26 Jul 2022 03:40:41 +0000 Subject: [PATCH] More implementation --- .../ApiService/Functions/Scaleset.cs | 40 ++++++++++++++++-- .../ApiService/OneFuzzTypes/Enums.cs | 40 +++++++++++++----- .../ApiService/OneFuzzTypes/Responses.cs | 5 +-- .../ApiService/onefuzzlib/NodeOperations.cs | 10 ++--- .../onefuzzlib/ScalesetOperations.cs | 42 +++++++++++++++---- src/ApiService/Tests/OrmModelsTest.cs | 5 +-- 6 files changed, 111 insertions(+), 31 deletions(-) diff --git a/src/ApiService/ApiService/Functions/Scaleset.cs b/src/ApiService/ApiService/Functions/Scaleset.cs index cc9fba97fcb..3c8e2754c4a 100644 --- a/src/ApiService/ApiService/Functions/Scaleset.cs +++ b/src/ApiService/ApiService/Functions/Scaleset.cs @@ -61,7 +61,41 @@ private async Task Patch(HttpRequestData req) { if (!request.IsOk) { return await _context.RequestHandling.NotOk(req, request.ErrorV, "ScalesetUpdate"); } - throw new NotImplementedException(); + + var answer = await _auth.CheckRequireAdmins(req); + if (!answer.IsOk) { + return await _context.RequestHandling.NotOk(req, answer.ErrorV, "ScalesetUpdate"); + } + + var scalesetResult = await _context.ScalesetOperations.GetById(request.OkV.ScalesetId); + if (!scalesetResult.IsOk) { + return await _context.RequestHandling.NotOk(req, scalesetResult.ErrorV, "ScalesetUpdate"); + } + + var scaleset = scalesetResult.OkV; + if (!scaleset.State.CanUpdate()) { + return await _context.RequestHandling.NotOk( + req, + new Error( + Code: ErrorCode.INVALID_REQUEST, + Errors: new[] { $"scaleset must be in one of the following states to update: {string.Join(", ", ScalesetStateHelper.CanUpdateStates)}" }), + "ScalesetUpdate"); + } + + if (request.OkV.Size is long size) { + var resizedScaleset = await _context.ScalesetOperations.SetSize(scaleset, size); + if (resizedScaleset.IsOk) { + scaleset = resizedScaleset.OkV; + } else { + return await _context.RequestHandling.NotOk( + req, + resizedScaleset.ErrorV, + "ScalesetUpdate"); + } + } + + scaleset = scaleset with { Auth = null }; + return await RequestHandling.Ok(req, ScalesetResponse.ForScaleset(scaleset)); } private async Task Get(HttpRequestData req) { @@ -79,7 +113,7 @@ private async Task Get(HttpRequestData req) { var scaleset = scalesetResult.OkV; - var response = ScalesetSearchResponse.ForScaleset(scaleset); + var response = ScalesetResponse.ForScaleset(scaleset); response = response with { Nodes = await _context.ScalesetOperations.GetNodes(scaleset) }; if (!search.IncludeAuth) { response = response with { Auth = null }; @@ -91,7 +125,7 @@ private async Task Get(HttpRequestData req) { var states = search.State ?? Enumerable.Empty(); var scalesets = await _context.ScalesetOperations.SearchStates(states).ToListAsync(); // don't return auths during list actions, only 'get' - var result = scalesets.Select(ss => ScalesetSearchResponse.ForScaleset(ss with { Auth = null })); + var result = scalesets.Select(ss => ScalesetResponse.ForScaleset(ss with { Auth = null })); return await RequestHandling.Ok(req, result); } } diff --git a/src/ApiService/ApiService/OneFuzzTypes/Enums.cs b/src/ApiService/ApiService/OneFuzzTypes/Enums.cs index 69e860e3339..0316bd02291 100644 --- a/src/ApiService/ApiService/OneFuzzTypes/Enums.cs +++ b/src/ApiService/ApiService/OneFuzzTypes/Enums.cs @@ -138,29 +138,49 @@ public static class JobStateHelper { public static class ScalesetStateHelper { - private static readonly IReadOnlySet _canUpdate = new HashSet { ScalesetState.Init, ScalesetState.Resize }; - private static readonly IReadOnlySet _needsWork = - new HashSet{ + private static readonly HashSet _canUpdate = + new() { + ScalesetState.Init, + ScalesetState.Resize, + }; + + private static readonly HashSet _needsWork = + new() { ScalesetState.Init, ScalesetState.Setup, ScalesetState.Resize, ScalesetState.Shutdown, - ScalesetState.Halt + ScalesetState.Halt, + }; + + private static readonly HashSet _available = + new() { + ScalesetState.Resize, + ScalesetState.Running, + }; + + private static readonly HashSet _resizing = + new() { + ScalesetState.Halt, + ScalesetState.Init, + ScalesetState.Setup, }; - private static readonly IReadOnlySet _available = new HashSet { ScalesetState.Resize, ScalesetState.Running }; - private static readonly IReadOnlySet _resizing = new HashSet { ScalesetState.Halt, ScalesetState.Init, ScalesetState.Setup }; /// set of states that indicate the scaleset can be updated - public static IReadOnlySet CanUpdate => _canUpdate; + public static bool CanUpdate(this ScalesetState state) => _canUpdate.Contains(state); + public static IReadOnlySet CanUpdateStates => _canUpdate; /// set of states that indicate work is needed during eventing - public static IReadOnlySet NeedsWork => _needsWork; + public static bool NeedsWork(this ScalesetState state) => _needsWork.Contains(state); + public static IReadOnlySet NeedsWorkStates => _needsWork; /// set of states that indicate if it's available for work - public static IReadOnlySet Available => _available; + public static bool IsAvailable(this ScalesetState state) => _available.Contains(state); + public static IReadOnlySet AvailableStates => _available; /// set of states that indicate scaleset is resizing - public static IReadOnlySet Resizing => _resizing; + public static bool IsResizing(this ScalesetState state) => _resizing.Contains(state); + public static IReadOnlySet ResizingStates => _resizing; } diff --git a/src/ApiService/ApiService/OneFuzzTypes/Responses.cs b/src/ApiService/ApiService/OneFuzzTypes/Responses.cs index 6bdbcf69f3b..a3e2b2ff746 100644 --- a/src/ApiService/ApiService/OneFuzzTypes/Responses.cs +++ b/src/ApiService/ApiService/OneFuzzTypes/Responses.cs @@ -91,7 +91,7 @@ public static JobResponse ForJob(Job j) ); } -public record ScalesetSearchResponse( +public record ScalesetResponse( PoolName PoolName, Guid ScalesetId, ScalesetState State, @@ -109,7 +109,7 @@ public record ScalesetSearchResponse( Dictionary Tags, List? Nodes ) : BaseResponse() { - public static ScalesetSearchResponse ForScaleset(Scaleset s) + public static ScalesetResponse ForScaleset(Scaleset s) => new( PoolName: s.PoolName, ScalesetId: s.ScalesetId, @@ -154,4 +154,3 @@ VmState State public record ProxyList( List Proxies ); - diff --git a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs index b6bdcadfc75..db067c051ff 100644 --- a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs @@ -133,25 +133,25 @@ public async Task CanProcessNewWork(Node node) { if (node.ScalesetId != null) { var scalesetResult = await _context.ScalesetOperations.GetById(node.ScalesetId.Value); - if (!scalesetResult.IsOk || scalesetResult.OkV == null) { + if (!scalesetResult.IsOk) { _logTracer.Info($"can_process_new_work invalid scaleset. scaleset_id:{node.ScalesetId} machine_id:{node.MachineId}"); return false; } - var scaleset = scalesetResult.OkV!; - if (!ScalesetStateHelper.Available.Contains(scaleset.State)) { + var scaleset = scalesetResult.OkV; + if (!scaleset.State.IsAvailable()) { _logTracer.Info($"can_process_new_work scaleset not available for work. scaleset_id:{node.ScalesetId} machine_id:{node.MachineId}"); return false; } } var poolResult = await _context.PoolOperations.GetByName(node.PoolName); - if (!poolResult.IsOk || poolResult.OkV == null) { + if (!poolResult.IsOk) { _logTracer.Info($"can_schedule - invalid pool. pool_name:{node.PoolName} machine_id:{node.MachineId}"); return false; } - var pool = poolResult.OkV!; + var pool = poolResult.OkV; if (!PoolStateHelper.Available.Contains(pool.State)) { _logTracer.Info($"can_schedule - pool is not available for work. pool_name:{node.PoolName} machine_id:{node.MachineId}"); return false; diff --git a/src/ApiService/ApiService/onefuzzlib/ScalesetOperations.cs b/src/ApiService/ApiService/onefuzzlib/ScalesetOperations.cs index 264c628448b..baa2d98a61a 100644 --- a/src/ApiService/ApiService/onefuzzlib/ScalesetOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/ScalesetOperations.cs @@ -15,6 +15,7 @@ public interface IScalesetOperations : IOrm { public Async.Task> GetNodes(Scaleset scaleset); IAsyncEnumerable SearchStates(IEnumerable states); Async.Task SetShutdown(Scaleset scaleset, bool now); + Async.Task> SetSize(Scaleset scaleset, long size); } public class ScalesetOperations : StatefulOrm, IScalesetOperations { @@ -37,17 +38,23 @@ public IAsyncEnumerable SearchByPool(PoolName poolName) { } - async Async.Task SetState(Scaleset scaleSet, ScalesetState state) { - if (scaleSet.State == state) - return; + async Async.Task> SetState(Scaleset scaleSet, ScalesetState state) { + if (scaleSet.State == state) { + return OneFuzzResult.Ok(scaleSet); + } - if (scaleSet.State == ScalesetState.Halt) - return; + if (scaleSet.State == ScalesetState.Halt) { + // terminal state, unable to change + // TODO: should this throw an exception instead? + return OneFuzzResult.Ok(scaleSet); + } var updatedScaleSet = scaleSet with { State = state }; - var r = await this.Replace(updatedScaleSet); + var r = await Update(updatedScaleSet); if (!r.IsOk) { - _log.Error($"Failed to update scaleset {scaleSet.ScalesetId} when updating state from {scaleSet.State} to {state}"); + var msg = "Failed to update scaleset {scaleSet.ScalesetId} when updating state from {scaleSet.State} to {state}"; + _log.Error(msg); + return OneFuzzResult.Error(ErrorCode.UNABLE_TO_UPDATE, msg); } if (state == ScalesetState.Resize) { @@ -59,6 +66,8 @@ await _context.Events.SendEvent( new EventScalesetStateUpdated(updatedScaleSet.ScalesetId, updatedScaleSet.PoolName, updatedScaleSet.State) ); } + + return OneFuzzResult.Ok(scaleSet); } async Async.Task SetFailed(Scaleset scaleSet, Error error) { @@ -365,4 +374,23 @@ public IAsyncEnumerable SearchStates(IEnumerable states public Async.Task SetShutdown(Scaleset scaleset, bool now) => SetState(scaleset, now ? ScalesetState.Halt : ScalesetState.Shutdown); + + public Async.Task> SetSize(Scaleset scaleset, long size) { + var permittedSize = Math.Min(size, MaxSize(scaleset)); + if (permittedSize == scaleset.Size) { + return Async.Task.FromResult(OneFuzzResult.Ok(scaleset)); // nothing to do + } + + scaleset = scaleset with { Size = permittedSize }; + return SetState(scaleset, ScalesetState.Resize); + } + + private static long MaxSize(Scaleset scaleset) { + // https://docs.microsoft.com/en-us/azure/virtual-machine-scale-sets/virtual-machine-scale-sets-placement-groups#checklist-for-using-large-scale-sets + if (scaleset.Image.StartsWith("/", StringComparison.Ordinal)) { + return 600; + } else { + return 1000; + } + } } diff --git a/src/ApiService/Tests/OrmModelsTest.cs b/src/ApiService/Tests/OrmModelsTest.cs index d713c82cb6a..6222adc10a7 100644 --- a/src/ApiService/Tests/OrmModelsTest.cs +++ b/src/ApiService/Tests/OrmModelsTest.cs @@ -220,7 +220,7 @@ public static Gen Task() { public static Gen Scaleset { get; } = from arg in Arb.Generate, - Tuple, Guid?>, + Tuple, Tuple>>>() from poolName in PoolNameGen select new Scaleset( @@ -237,8 +237,7 @@ from poolName in PoolNameGen EphemeralOsDisks: arg.Item2.Item3, NeedsConfigUpdate: arg.Item2.Item4, Error: arg.Item2.Item5, - Nodes: arg.Item2.Item6, - ClientId: arg.Item2.Item7, + ClientId: arg.Item2.Item6, ClientObjectId: arg.Item3.Item1, Tags: arg.Item3.Item2);