Skip to content
This repository was archived by the owner on Nov 1, 2023. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/ApiService/ApiService/Functions/AgentEvents.cs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,10 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
}

if (ev.State == NodeState.Free) {
if (!node.Managed) {
return null;
}

if (node.ReimageRequested || node.DeleteRequested) {
_log.Info($"stopping free node with reset flags: {machineId:Tag:MachineId}");
// discard result: node not used after this point
Expand Down
39 changes: 38 additions & 1 deletion src/ApiService/ApiService/Functions/Pool.cs
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@ public Pool(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context
}

[Function("Pool")]
public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req)
public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE", "PATCH")] HttpRequestData req)
=> _auth.CallIfUser(req, r => r.Method switch {
"GET" => Get(r),
"POST" => Post(r),
"DELETE" => Delete(r),
"PATCH" => Patch(r),
var m => throw new InvalidOperationException("Unsupported HTTP method {m}"),
});

Expand Down Expand Up @@ -71,6 +72,42 @@ private async Task<HttpResponseData> Post(HttpRequestData req) {
return await RequestHandling.Ok(req, await Populate(PoolToPoolResponse(newPool), true));
}


private async Task<HttpResponseData> Patch(HttpRequestData req) {
var request = await RequestHandling.ParseRequest<PoolUpdate>(req);
if (!request.IsOk) {
return await _context.RequestHandling.NotOk(req, request.ErrorV, "PoolUpdate");
}

var answer = await _auth.CheckRequireAdmins(req);
if (!answer.IsOk) {
return await _context.RequestHandling.NotOk(req, answer.ErrorV, "PoolUpdate");
}

var update = request.OkV;
var pool = await _context.PoolOperations.GetByName(update.Name);
if (!pool.IsOk) {
return await _context.RequestHandling.NotOk(
req,
new Error(
Code: ErrorCode.INVALID_REQUEST,
Errors: new string[] { "pool with that name does not exist exists" }),
"PoolUpdate");
}

var updated = pool.OkV with { ObjectId = update.ObjectId };
var updatePool = await _context.PoolOperations.Update(updated);
if (updatePool.IsOk) {
return await RequestHandling.Ok(req, await Populate(PoolToPoolResponse(updated), true));
} else {
return await _context.RequestHandling.NotOk(req, new Error(
Code: ErrorCode.INVALID_REQUEST,
Errors: new string[] { updatePool.ErrorV.Reason }), "PoolUpdate");
}


}

private async Task<HttpResponseData> Get(HttpRequestData req) {
var request = await RequestHandling.ParseRequest<PoolSearch>(req);
if (!request.IsOk) {
Expand Down
7 changes: 7 additions & 0 deletions src/ApiService/ApiService/OneFuzzTypes/Requests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,13 @@ public record PoolCreate(
[property: Required] Os Os,
[property: Required] Architecture Arch,
[property: Required] bool Managed,
Guid? ObjectId = null,
bool Update = false
) : BaseRequest;


public record PoolUpdate(
[property: Required] PoolName Name,
Guid? ObjectId = null
) : BaseRequest;

Expand Down
7 changes: 4 additions & 3 deletions src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
Original file line number Diff line number Diff line change
Expand Up @@ -195,12 +195,12 @@ public async Task<CanProcessNewWorkResponse> CanProcessNewWork(Node node) {
return CanProcessNewWorkResponse.NotAllowed("node is set to be deleted");
}

if (node.ReimageRequested) {
if (node.ReimageRequested && node.Managed) {
_ = await Stop(node, done: true);
return CanProcessNewWorkResponse.NotAllowed("node is set to be reimaged");
}

if (await CouldShrinkScaleset(node)) {
if (await CouldShrinkScaleset(node) && node.Managed) {
_ = await SetHalt(node);
return CanProcessNewWorkResponse.NotAllowed("node is scheduled to shrink");
}
Expand Down Expand Up @@ -488,7 +488,8 @@ public bool IsOutdated(Node node) {
}

public bool IsTooOld(Node node) {
return node.ScalesetId != null
return node.Managed
&& node.ScalesetId != null
&& node.InitializedAt != null
&& node.InitializedAt < DateTime.UtcNow - INodeOperations.NODE_REIMAGE_TIME;
}
Expand Down
18 changes: 18 additions & 0 deletions src/cli/onefuzz/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -1297,6 +1297,24 @@ def create(
),
)

def update(
self,
name: str,
object_id: Optional[UUID] = None,
) -> models.Pool:
"""
Update a worker pool

:param str name: Name of the worker-pool
"""
self.logger.debug("create worker pool")

return self._req_model(
"PATCH",
models.Pool,
data=requests.PoolUpdate(name=name, object_id=object_id),
)

def get_config(self, pool_name: primitives.PoolName) -> models.AgentConfig:
"""Get the agent configuration for the pool"""

Expand Down
1 change: 1 addition & 0 deletions src/pytypes/onefuzztypes/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -652,6 +652,7 @@ class Pool(BaseModel):
# intended to be used to pass the information to the CLI when the CLI asks
# for information about what work is in the queue for the pool.
work_queue: Optional[List[WorkSetSummary]]
object_id: Optional[UUID]

# explicitly excluded from Tables
scaleset_summary: Optional[List[ScalesetSummary]]
Expand Down
5 changes: 5 additions & 0 deletions src/pytypes/onefuzztypes/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,11 @@ class PoolCreate(BaseRequest):
autoscale: Optional[AutoScaleConfig]


class PoolUpdate(BaseRequest):
name: PoolName
object_id: Optional[UUID]


class PoolSearch(BaseRequest):
pool_id: Optional[UUID]
name: Optional[PoolName]
Expand Down