From 7fbcbe5b1780b6fc576256e7bbfb80fa191e34af Mon Sep 17 00:00:00 2001
From: Tim Holm <tim.holm@nitric.io>
Date: Fri, 18 Oct 2024 15:04:10 +1100
Subject: [PATCH] fix(aws/batch): Add config for increasing docker system
 volume size via launch template (#686)

* add launch template to increase default EBS size.
* add block device mapping to launch template config
---
 cloud/aws/common/config.go | 25 +++++++++++++++++++------
 cloud/aws/deploy/batch.go  | 25 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/cloud/aws/common/config.go b/cloud/aws/common/config.go
index 5cf830f38..634f1ff13 100644
--- a/cloud/aws/common/config.go
+++ b/cloud/aws/common/config.go
@@ -33,10 +33,22 @@ type AwsImports struct {
 	Buckets map[string]string
 }
 
+type EcsLaunchTemplate struct {
+	BlockDeviceMappings []struct {
+		DeviceName string `mapstructure:"device-name,omitempty"`
+		Ebs        struct {
+			DeleteOnTermination string `mapstructure:"delete-on-termination,omitempty"`
+			VolumeSize          int    `mapstructure:"volume-size,omitempty"`
+			VolumeType          string `mapstructure:"volume-type,omitempty"`
+		} `mapstructure:"ebs,omitempty"`
+	} `mapstructure:"block-device-mappings,omitempty"`
+}
+
 type BatchComputeEnvConfig struct {
-	MinCpus       int      `mapstructure:"min-cpus"`
-	MaxCpus       int      `mapstructure:"max-cpus"`
-	InstanceTypes []string `mapstructure:"instance-types"`
+	MinCpus        int                `mapstructure:"min-cpus"`
+	MaxCpus        int                `mapstructure:"max-cpus"`
+	InstanceTypes  []string           `mapstructure:"instance-types"`
+	LaunchTemplate *EcsLaunchTemplate `mapstructure:"launch-template,omitempty"`
 }
 
 type AwsConfig struct {
@@ -72,9 +84,10 @@ var defaultLambdaConfig = &AwsLambdaConfig{
 }
 
 var defaultBatchComputeEnvConfig = &BatchComputeEnvConfig{
-	MinCpus:       0,
-	MaxCpus:       32,
-	InstanceTypes: []string{"optimal"},
+	MinCpus:        0,
+	MaxCpus:        32,
+	InstanceTypes:  []string{"optimal"},
+	LaunchTemplate: nil,
 }
 
 var defaultAwsConfigItem = AwsConfigItem{
diff --git a/cloud/aws/deploy/batch.go b/cloud/aws/deploy/batch.go
index caaefdbe3..a47b2d82c 100644
--- a/cloud/aws/deploy/batch.go
+++ b/cloud/aws/deploy/batch.go
@@ -149,6 +149,31 @@ func (a *NitricAwsPulumiProvider) batch(ctx *pulumi.Context) error {
 		InstanceRole:     instanceProfile.Arn,
 	}
 
+	if a.AwsConfig.BatchComputeEnvConfig.LaunchTemplate != nil {
+		blockDevicemappings := awsec2.LaunchTemplateBlockDeviceMappingArray{}
+		for _, bd := range a.AwsConfig.BatchComputeEnvConfig.LaunchTemplate.BlockDeviceMappings {
+			blockDevicemappings = append(blockDevicemappings, &awsec2.LaunchTemplateBlockDeviceMappingArgs{
+				DeviceName: pulumi.String(bd.DeviceName),
+				Ebs: &awsec2.LaunchTemplateBlockDeviceMappingEbsArgs{
+					DeleteOnTermination: pulumi.String(bd.Ebs.DeleteOnTermination),
+					VolumeSize:          pulumi.Int(bd.Ebs.VolumeSize),
+					VolumeType:          pulumi.String(bd.Ebs.VolumeType),
+				},
+			})
+		}
+
+		launchTemplate, err := awsec2.NewLaunchTemplate(ctx, "batch-launch-template", &awsec2.LaunchTemplateArgs{
+			BlockDeviceMappings: blockDevicemappings,
+		})
+		if err != nil {
+			return err
+		}
+
+		computeResourceOptions.LaunchTemplate = &batch.ComputeEnvironmentComputeResourcesLaunchTemplateArgs{
+			LaunchTemplateName: launchTemplate.Name,
+		}
+	}
+
 	a.ComputeEnvironment, err = batch.NewComputeEnvironment(ctx, "compute-environment", &batch.ComputeEnvironmentArgs{
 		ComputeEnvironmentName: pulumi.Sprintf("%s-compute-environment", a.StackId),
 		ComputeResources:       computeResourceOptions,