Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Cache VMSS VM InstanceID lookups #2464

Merged
merged 5 commits into from
Sep 30, 2022
Merged

Cache VMSS VM InstanceID lookups #2464

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7440b73
Cache VMSS VM InstanceID lookups
Porges Sep 29, 2022
bb2135a
Adding an expiration time to the cache
chkeita Sep 29, 2022
349263d
make the TTL 10 min
chkeita Sep 29, 2022
1c4c970
properly add entries to the cache
chkeita Sep 29, 2022
12f67f7
Merge branch 'main' into cache-instanceid-lookup
chkeita Sep 30, 2022
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
3 changes: 0 additions & 3 deletions src/ApiService/ApiService/Functions/AgentCanSchedule.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
var canScheduleRequest = request.OkV;

var node = await _context.NodeOperations.GetByMachineId(canScheduleRequest.MachineId);

if (node == null) {
_log.Warning($"Unable to find {canScheduleRequest.MachineId:Tag:MachineId}");
return await _context.RequestHandling.NotOk(
Expand All @@ -44,14 +43,12 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
}

var allowed = true;

if (!await _context.NodeOperations.CanProcessNewWork(node)) {
allowed = false;
}

var task = await _context.TaskOperations.GetByTaskId(canScheduleRequest.TaskId);
var workStopped = task == null || task.State.ShuttingDown();

if (workStopped) {
_log.Info($"Work stopped for: {canScheduleRequest.MachineId:Tag:MachineId} and {canScheduleRequest.TaskId:Tag:TaskId}");
allowed = false;
Expand Down
8 changes: 4 additions & 4 deletions src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,17 +103,17 @@ await ScalesetNodeExists(node)) {
}

public async Async.Task<bool> ScalesetNodeExists(Node node) {
if (node.ScalesetId == null) {
var scalesetId = node.ScalesetId;
if (scalesetId is null) {
return false;
}

var scalesetResult = await _context.ScalesetOperations.GetById((Guid)(node.ScalesetId!));
var scalesetResult = await _context.ScalesetOperations.GetById(scalesetId.Value);
if (!scalesetResult.IsOk || scalesetResult.OkV == null) {
return false;
}
var scaleset = scalesetResult.OkV;

var instanceId = await _context.VmssOperations.GetInstanceId(scaleset.ScalesetId, node.MachineId);
var instanceId = await _context.VmssOperations.GetInstanceId(scalesetResult.OkV.ScalesetId, node.MachineId);
return instanceId.IsOk;
}

Expand Down
101 changes: 65 additions & 36 deletions src/ApiService/ApiService/onefuzzlib/VmssOperations.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,21 @@ public async Async.Task<OneFuzzResultVoid> ResizeVmss(Guid name, long capacity)


private VirtualMachineScaleSetResource GetVmssResource(Guid name) {
var resourceGroup = _creds.GetBaseResourceGroup();
var id = VirtualMachineScaleSetResource.CreateResourceIdentifier(_creds.GetSubscription(), resourceGroup, name.ToString());
var id = VirtualMachineScaleSetResource.CreateResourceIdentifier(
_creds.GetSubscription(),
_creds.GetBaseResourceGroup(),
name.ToString());
return _creds.ArmClient.GetVirtualMachineScaleSetResource(id);
}

private VirtualMachineScaleSetVmResource GetVmssVmResource(Guid name, string instanceId) {
var id = VirtualMachineScaleSetVmResource.CreateResourceIdentifier(
_creds.GetSubscription(),
_creds.GetBaseResourceGroup(),
name.ToString(),
instanceId);
return _creds.ArmClient.GetVirtualMachineScaleSetVmResource(id);
}

public async Async.Task<VirtualMachineScaleSetData?> GetVmss(Guid name) {
try {
Expand Down Expand Up @@ -180,58 +190,77 @@ public async Async.Task<IDictionary<Guid, string>> ListInstanceIds(Guid name) {
return results;
}

private record InstanceIdKey(Guid Scaleset, Guid VmId);
private Task<string> GetInstanceIdForVmId(Guid scaleset, Guid vmId)
=> _cache.GetOrCreateAsync(new InstanceIdKey(scaleset, vmId), async entry => {
var scalesetResource = GetVmssResource(scaleset);
var vmIdString = vmId.ToString();
await foreach (var vm in scalesetResource.GetVirtualMachineScaleSetVms().AsAsyncEnumerable()) {
var response = await vm.GetAsync();
var instanceId = response.Value.Data.InstanceId;
if (response.Value.Data.VmId == vmIdString) {
// we found the VM we are looking for
entry.AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(10);
return instanceId;
} else {
// if we find any other VMs, put them in the cache
if (Guid.TryParse(response.Value.Data.VmId, out var vmId)) {
using var e = _cache.CreateEntry(new InstanceIdKey(scaleset, vmId));
_ = e.SetValue(instanceId);
e.AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(10);
}
}
}

throw new Exception($"unable to find instance ID for scaleset vm {scaleset}:{vmId}");
});

public async Async.Task<OneFuzzResult<VirtualMachineScaleSetVmResource>> GetInstanceVm(Guid name, Guid vmId) {
_log.Info($"get instance ID for scaleset node: {name:Tag:VmssName}:{vmId:Tag:VmId}");
var scaleSet = GetVmssResource(name);
var instanceId = await GetInstanceId(name, vmId);
if (!instanceId.IsOk) {
return instanceId.ErrorV;
}

var resource = GetVmssVmResource(name, instanceId.OkV);
try {
await foreach (var vm in scaleSet.GetVirtualMachineScaleSetVms().AsAsyncEnumerable()) {
var response = await vm.GetAsync();
if (!response.Value.HasData) {
return OneFuzzResult<VirtualMachineScaleSetVmResource>.Error(ErrorCode.UNABLE_TO_FIND, $"failed to get vm data");
}

if (response.Value.Data.VmId == vmId.ToString()) {
return OneFuzzResult<VirtualMachineScaleSetVmResource>.Ok(response);
}
}
var response = await resource.GetAsync();
return OneFuzzResult.Ok(response.Value);
} catch (Exception ex) when (ex is RequestFailedException || ex is CloudException) {
_log.Exception(ex, $"unable to find vm instance: {name:Tag:VmssName}:{vmId:Tag:VmId}");
return OneFuzzResult<VirtualMachineScaleSetVmResource>.Error(ErrorCode.UNABLE_TO_FIND, $"unable to find vm instance: {name}:{vmId}");
return OneFuzzResult<VirtualMachineScaleSetVmResource>.Error(ErrorCode.UNABLE_TO_FIND, $"unable to find vm instance: {name}:{instanceId}");
}
return OneFuzzResult<VirtualMachineScaleSetVmResource>.Error(ErrorCode.UNABLE_TO_FIND, $"unable to find scaleset machine: {name}:{vmId}");
}

public async Async.Task<OneFuzzResult<string>> GetInstanceId(Guid name, Guid vmId) {
var vm = await GetInstanceVm(name, vmId);
if (vm.IsOk) {
return OneFuzzResult<string>.Ok(vm.OkV!.Data.InstanceId);
} else {
return OneFuzzResult<string>.Error(vm.ErrorV);
try {
return OneFuzzResult.Ok(await GetInstanceIdForVmId(name, vmId));
} catch {
return new Error(ErrorCode.UNABLE_TO_FIND, new string[] { $"unable to find scaleset machine: {name}:{vmId}" });
}
}

public async Async.Task<OneFuzzResultVoid> UpdateScaleInProtection(Guid name, Guid vmId, bool protectFromScaleIn) {
var res = await GetInstanceVm(name, vmId);
if (!res.IsOk) {
return OneFuzzResultVoid.Error(res.ErrorV);
} else {
var instanceVm = res.OkV;
instanceVm.Data.ProtectionPolicy ??= new();
if (instanceVm.Data.ProtectionPolicy.ProtectFromScaleIn != protectFromScaleIn) {
instanceVm.Data.ProtectionPolicy.ProtectFromScaleIn = protectFromScaleIn;
var vmCollection = GetVmssResource(name).GetVirtualMachineScaleSetVms();
try {
_ = await vmCollection.CreateOrUpdateAsync(WaitUntil.Started, instanceVm.Data.InstanceId, instanceVm.Data);
return OneFuzzResultVoid.Ok;
} catch {
var msg = $"unable to set protection policy on: {vmId}:{instanceVm.Id} in vmss {name}";
return OneFuzzResultVoid.Error(ErrorCode.UNABLE_TO_UPDATE, msg);
}
} else {
_log.Info($"scale in protection was already set to {protectFromScaleIn:Tag:ProtectFromScaleIn} on vm {vmId:Tag:VmId} for scaleset {name:Tag:VmssName}");
return res.ErrorV;
}

var instanceVm = res.OkV;
instanceVm.Data.ProtectionPolicy ??= new();
if (instanceVm.Data.ProtectionPolicy.ProtectFromScaleIn != protectFromScaleIn) {
instanceVm.Data.ProtectionPolicy.ProtectFromScaleIn = protectFromScaleIn;
var vmCollection = GetVmssResource(name).GetVirtualMachineScaleSetVms();
try {
_ = await vmCollection.CreateOrUpdateAsync(WaitUntil.Started, instanceVm.Data.InstanceId, instanceVm.Data);
return OneFuzzResultVoid.Ok;
} catch {
var msg = $"unable to set protection policy on: {vmId}:{instanceVm.Id} in vmss {name}";
return OneFuzzResultVoid.Error(ErrorCode.UNABLE_TO_UPDATE, msg);
}
} else {
_log.Info($"scale in protection was already set to {protectFromScaleIn:Tag:ProtectFromScaleIn} on vm {vmId:Tag:VmId} for scaleset {name:Tag:VmssName}");
return OneFuzzResultVoid.Ok;
}
}

Expand Down