Skip to content

Commit 7e3b431

Browse files
bluna301bluna301
committed
shm_size for Docker and Podman plugins
Signed-off-by: bluna301 <luna.bryanr@gmail.com>
1 parent 59f46e9 commit 7e3b431

9 files changed

Lines changed: 204 additions & 0 deletions

File tree

‎src/TaskManager/Plug-ins/Docker/DockerPlugin.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,18 @@ private CreateContainerParameters BuildContainerSpecification(IList<ContainerVol
367367
},
368368
};
369369

370+
if (Event.TaskPluginArguments.ContainsKey(Keys.ShmSize))
371+
{
372+
if (long.TryParse(Event.TaskPluginArguments[Keys.ShmSize], out var shmsize))
373+
{
374+
parameters.HostConfig.ShmSize = shmsize;
375+
}
376+
else
377+
{
378+
_logger.InvalidShmSize(Event.TaskPluginArguments[Keys.ShmSize]);
379+
}
380+
}
381+
370382
if (Event.TaskPluginArguments.ContainsKey(Keys.User))
371383
{
372384
parameters.User = Event.TaskPluginArguments[Keys.User];

‎src/TaskManager/Plug-ins/Docker/Keys.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,11 @@ internal static class Keys
5858
/// </summary>
5959
public static readonly string TemporaryStorageContainerPath = "temp_storage_container_path";
6060

61+
/// <summary>
62+
/// Key for setting the shm size.
63+
/// </summary>
64+
public static readonly string ShmSize = "shm_size";
65+
6166
/// <summary>
6267
/// Prefix for envrionment variables.
6368
/// </summary>

‎src/TaskManager/Plug-ins/Docker/Logging/Log.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,5 +115,8 @@ public static partial class Log
115115

116116
[LoggerMessage(EventId = 1031, Level = LogLevel.Error, Message = "Error setting directory {path} with permission {user}.")]
117117
public static partial void ErrorSettingDirectoryPermission(this ILogger logger, string path, string user);
118+
119+
[LoggerMessage(EventId = 1032, Level = LogLevel.Error, Message = "Invalid size specified for /dev/shm: {size}. Please use value between 1 and 9223372036854775807.")]
120+
public static partial void InvalidShmSize(this ILogger logger, string size);
118121
}
119122
}

‎src/TaskManager/Plug-ins/Podman/IPodmanContainerCreator.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,9 @@ public class PodmanCreateContainerRequest
107107

108108
[JsonPropertyName("user")]
109109
public string? User { get; set; }
110+
111+
[JsonPropertyName("shm_size")]
112+
public long? ShmSize { get; set; }
110113
}
111114

112115
public class PodmanDevice

‎src/TaskManager/Plug-ins/Podman/Keys.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,11 @@ internal static class Keys
5858
/// </summary>
5959
public static readonly string TemporaryStorageContainerPath = "temp_storage_container_path";
6060

61+
/// <summary>
62+
/// Key for setting the shm size.
63+
/// </summary>
64+
public static readonly string ShmSize = "shm_size";
65+
6166
/// <summary>
6267
/// Prefix for environment variables.
6368
/// </summary>

‎src/TaskManager/Plug-ins/Podman/Logging/Log.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,5 +115,8 @@ public static partial class Log
115115

116116
[LoggerMessage(EventId = 1031, Level = LogLevel.Error, Message = "Error setting directory {path} with permission {user}.")]
117117
public static partial void ErrorSettingDirectoryPermission(this ILogger logger, Exception exception, string path, string user);
118+
119+
[LoggerMessage(EventId = 1032, Level = LogLevel.Error, Message = "Invalid size specified for /dev/shm: {size}. Please use value between 1 and 9223372036854775807.")]
120+
public static partial void InvalidShmSize(this ILogger logger, string size);
118121
}
119122
}

‎src/TaskManager/Plug-ins/Podman/PodmanPlugin.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,18 @@ private PodmanCreateContainerRequest BuildContainerSpecification(IList<Container
422422
request.User = Event.TaskPluginArguments[Keys.User];
423423
}
424424

425+
if (Event.TaskPluginArguments.ContainsKey(Keys.ShmSize))
426+
{
427+
if (long.TryParse(Event.TaskPluginArguments[Keys.ShmSize], out var shmsize))
428+
{
429+
request.ShmSize = shmsize;
430+
}
431+
else
432+
{
433+
_logger.InvalidShmSize(Event.TaskPluginArguments[Keys.ShmSize]);
434+
}
435+
}
436+
425437
return request;
426438
}
427439

‎tests/UnitTests/TaskManager.Docker.Tests/DockerPluginTest.cs‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,80 @@ public async Task GetStatus_WhenCalled_ExpectToWaitUntilFinalStatuses()
474474
It.IsAny<CancellationToken>()), Times.Exactly(3));
475475
}
476476

477+
[Fact(DisplayName = "ExecuteTask - when shm_size is specified expect it to be set on HostConfig")]
478+
public async Task ExecuteTask_WhenShmSizeIsSpecified_ExpectItToBeSetOnHostConfig()
479+
{
480+
var payloadFiles = new List<VirtualFileInfo>()
481+
{
482+
new VirtualFileInfo( "file.dcm", "path/to/file.dcm", "etag", 1000)
483+
};
484+
var contianerId = Guid.NewGuid().ToString();
485+
486+
_dockerClient.Setup(p => p.Images.CreateImageAsync(
487+
It.IsAny<ImagesCreateParameters>(),
488+
It.IsAny<AuthConfig>(),
489+
It.IsAny<IProgress<JSONMessage>>(),
490+
It.IsAny<CancellationToken>()));
491+
_dockerClient.Setup(p => p.Containers.CreateContainerAsync(
492+
It.IsAny<CreateContainerParameters>(),
493+
It.IsAny<CancellationToken>()))
494+
.ReturnsAsync(new CreateContainerResponse { ID = contianerId, Warnings = new List<string>() });
495+
496+
_storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
497+
.ReturnsAsync(payloadFiles);
498+
_storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
499+
.ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
500+
501+
var message = GenerateTaskDispatchEventWithValidArguments();
502+
message.TaskPluginArguments[Keys.ShmSize] = "2147483648";
503+
504+
var runner = new DockerPlugin(_serviceScopeFactory.Object, _logger.Object, message);
505+
var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext);
506+
507+
Assert.Equal(TaskExecutionStatus.Accepted, result.Status);
508+
_dockerClient.Verify(p => p.Containers.CreateContainerAsync(
509+
It.Is<CreateContainerParameters>(c => c.HostConfig.ShmSize == 2147483648),
510+
It.IsAny<CancellationToken>()), Times.Once());
511+
512+
runner.Dispose();
513+
}
514+
515+
[Fact(DisplayName = "ExecuteTask - when shm_size is invalid expect log warning and no exception")]
516+
public async Task ExecuteTask_WhenShmSizeIsInvalid_ExpectLogWarningAndNoException()
517+
{
518+
var payloadFiles = new List<VirtualFileInfo>()
519+
{
520+
new VirtualFileInfo( "file.dcm", "path/to/file.dcm", "etag", 1000)
521+
};
522+
var contianerId = Guid.NewGuid().ToString();
523+
524+
_dockerClient.Setup(p => p.Images.CreateImageAsync(
525+
It.IsAny<ImagesCreateParameters>(),
526+
It.IsAny<AuthConfig>(),
527+
It.IsAny<IProgress<JSONMessage>>(),
528+
It.IsAny<CancellationToken>()));
529+
_dockerClient.Setup(p => p.Containers.CreateContainerAsync(
530+
It.IsAny<CreateContainerParameters>(),
531+
It.IsAny<CancellationToken>()))
532+
.ReturnsAsync(new CreateContainerResponse { ID = contianerId, Warnings = new List<string>() });
533+
534+
_storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
535+
.ReturnsAsync(payloadFiles);
536+
_storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
537+
.ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
538+
539+
var message = GenerateTaskDispatchEventWithValidArguments();
540+
message.TaskPluginArguments[Keys.ShmSize] = "not-a-number";
541+
542+
var runner = new DockerPlugin(_serviceScopeFactory.Object, _logger.Object, message);
543+
var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext);
544+
545+
Assert.Equal(TaskExecutionStatus.Accepted, result.Status);
546+
_logger.VerifyLogging("Invalid size specified for /dev/shm: not-a-number. Please use value between 1 and 9223372036854775807.", LogLevel.Error, Times.Once());
547+
548+
runner.Dispose();
549+
}
550+
477551
[Fact(DisplayName = "HandleTimeout - when called expecte to terminate container")]
478552
public async Task HandleTimeout_WhenCalled_ExpectToTerminateContainer()
479553
{

‎tests/UnitTests/TaskManager.Podman.Tests/PodmanPluginTest.cs‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -532,6 +532,93 @@ public async Task GetStatus_WhenCalled_ExpectToWaitUntilFinalStatuses()
532532
It.IsAny<CancellationToken>()), Times.Exactly(3));
533533
}
534534

535+
[Fact(DisplayName = "ExecuteTask - when shm_size is specified expect it to be set on container request")]
536+
public async Task ExecuteTask_WhenShmSizeIsSpecified_ExpectItToBeSetOnContainerRequest()
537+
{
538+
var payloadFiles = new List<VirtualFileInfo>()
539+
{
540+
new VirtualFileInfo( "file.dcm", "path/to/file.dcm", "etag", 1000)
541+
};
542+
var contianerId = Guid.NewGuid().ToString();
543+
544+
_podmanClient.Setup(p => p.Images.CreateImageAsync(
545+
It.IsAny<ImagesCreateParameters>(),
546+
It.IsAny<AuthConfig>(),
547+
It.IsAny<IProgress<JSONMessage>>(),
548+
It.IsAny<CancellationToken>()));
549+
_containerCreator.Setup(p => p.CreateContainerAsync(
550+
It.IsAny<Uri>(),
551+
It.IsAny<PodmanCreateContainerRequest>(),
552+
It.IsAny<CancellationToken>()))
553+
.ReturnsAsync(new PodmanCreateContainerResponse { Id = contianerId, Warnings = new List<string>() });
554+
_podmanClient.Setup(p => p.Containers.StartContainerAsync(
555+
It.IsAny<string>(),
556+
It.IsAny<ContainerStartParameters>(),
557+
It.IsAny<CancellationToken>()))
558+
.ReturnsAsync(true);
559+
560+
_storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
561+
.ReturnsAsync(payloadFiles);
562+
_storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
563+
.ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
564+
565+
var message = GenerateTaskDispatchEventWithValidArguments();
566+
message.TaskPluginArguments[Keys.ShmSize] = "2147483648";
567+
568+
var runner = new PodmanPlugin(_serviceScopeFactory.Object, _logger.Object, message);
569+
var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext);
570+
571+
Assert.Equal(TaskExecutionStatus.Accepted, result.Status);
572+
_containerCreator.Verify(p => p.CreateContainerAsync(
573+
It.IsAny<Uri>(),
574+
It.Is<PodmanCreateContainerRequest>(r => r.ShmSize == 2147483648),
575+
It.IsAny<CancellationToken>()), Times.Once());
576+
577+
runner.Dispose();
578+
}
579+
580+
[Fact(DisplayName = "ExecuteTask - when shm_size is invalid expect log warning and no exception")]
581+
public async Task ExecuteTask_WhenShmSizeIsInvalid_ExpectLogWarningAndNoException()
582+
{
583+
var payloadFiles = new List<VirtualFileInfo>()
584+
{
585+
new VirtualFileInfo( "file.dcm", "path/to/file.dcm", "etag", 1000)
586+
};
587+
var contianerId = Guid.NewGuid().ToString();
588+
589+
_podmanClient.Setup(p => p.Images.CreateImageAsync(
590+
It.IsAny<ImagesCreateParameters>(),
591+
It.IsAny<AuthConfig>(),
592+
It.IsAny<IProgress<JSONMessage>>(),
593+
It.IsAny<CancellationToken>()));
594+
_containerCreator.Setup(p => p.CreateContainerAsync(
595+
It.IsAny<Uri>(),
596+
It.IsAny<PodmanCreateContainerRequest>(),
597+
It.IsAny<CancellationToken>()))
598+
.ReturnsAsync(new PodmanCreateContainerResponse { Id = contianerId, Warnings = new List<string>() });
599+
_podmanClient.Setup(p => p.Containers.StartContainerAsync(
600+
It.IsAny<string>(),
601+
It.IsAny<ContainerStartParameters>(),
602+
It.IsAny<CancellationToken>()))
603+
.ReturnsAsync(true);
604+
605+
_storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
606+
.ReturnsAsync(payloadFiles);
607+
_storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
608+
.ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
609+
610+
var message = GenerateTaskDispatchEventWithValidArguments();
611+
message.TaskPluginArguments[Keys.ShmSize] = "not-a-number";
612+
613+
var runner = new PodmanPlugin(_serviceScopeFactory.Object, _logger.Object, message);
614+
var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext);
615+
616+
Assert.Equal(TaskExecutionStatus.Accepted, result.Status);
617+
_logger.VerifyLogging("Invalid size specified for /dev/shm: not-a-number. Please use value between 1 and 9223372036854775807.", LogLevel.Error, Times.Once());
618+
619+
runner.Dispose();
620+
}
621+
535622
[Fact(DisplayName = "HandleTimeout - when called expecte to terminate container")]
536623
public async Task HandleTimeout_WhenCalled_ExpectToTerminateContainer()
537624
{

0 commit comments

Comments
 (0)