aws-infrastructure.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: 'Flight Server Infrastructure - EC2 instance behind Network Load Balancer for gRPC traffic'

Parameters:
  InstanceType:
    Type: String
    Default: t3.micro
    AllowedValues:
      - t3.micro
      - t3.small
      - t3.medium
      - t3.large
      - t3.xlarge
      - m5.large
      - m5.xlarge
    Description: EC2 instance type for the Flight server
  
  KeyPairName:
    Type: AWS::EC2::KeyPair::KeyName
    Description: Name of an existing EC2 KeyPair to enable SSH access to the instance
  
  AllowedCIDR:
    Type: String
    Default: 0.0.0.0/0
    Description: CIDR block allowed to access the Flight server (default allows all)
    AllowedPattern: ^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/([0-9]|[1-2][0-9]|3[0-2]))$

Resources:
  # VPC and Networking
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 10.0.0.0/16
      EnableDnsHostnames: true
      EnableDnsSupport: true
      Tags:
        - Key: Name
          Value: FlightServer-VPC

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: FlightServer-IGW

  AttachGateway:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref VPC
      InternetGatewayId: !Ref InternetGateway

  PublicSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      CidrBlock: 10.0.1.0/24
      AvailabilityZone: !Select [0, !GetAZs '']
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: FlightServer-PublicSubnet1

  PublicSubnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      CidrBlock: 10.0.2.0/24
      AvailabilityZone: !Select [1, !GetAZs '']
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: FlightServer-PublicSubnet2

  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: FlightServer-PublicRouteTable

  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: AttachGateway
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway

  PublicSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnet1
      RouteTableId: !Ref PublicRouteTable

  PublicSubnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnet2
      RouteTableId: !Ref PublicRouteTable

  # Security Groups
  FlightServerSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for Flight Server EC2 instance
      VpcId: !Ref VPC
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 8815
          ToPort: 8815
          SourceSecurityGroupId: !Ref LoadBalancerSecurityGroup
          Description: Flight server port from Load Balancer
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: !Ref AllowedCIDR
          Description: SSH access
      Tags:
        - Key: Name
          Value: FlightServer-EC2-SG

  LoadBalancerSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for Network Load Balancer
      VpcId: !Ref VPC
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 8815
          ToPort: 8815
          CidrIp: !Ref AllowedCIDR
          Description: gRPC Flight traffic
      Tags:
        - Key: Name
          Value: FlightServer-NLB-SG

  # IAM Role for EC2 instance
  FlightServerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: ec2.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/CloudWatchAgentServerPolicy
      Tags:
        - Key: Name
          Value: FlightServer-Role

  FlightServerInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
        - !Ref FlightServerRole

  # EC2 Instance
  FlightServerInstance:
    Type: AWS::EC2::Instance
    Properties:
      ImageId: ami-0c02fb55956c7d316  # Amazon Linux 2023 AMI (update for your region)
      InstanceType: !Ref InstanceType
      KeyName: !Ref KeyPairName
      SecurityGroupIds:
        - !Ref FlightServerSecurityGroup
      SubnetId: !Ref PublicSubnet1
      IamInstanceProfile: !Ref FlightServerInstanceProfile
      UserData:
        Fn::Base64: !Sub |
          #!/bin/bash
          yum update -y
          
          # Install Java 21
          yum install -y java-21-amazon-corretto-headless
          
          # Create flight user
          useradd -m -s /bin/bash flight
          
          # Create application directory
          mkdir -p /opt/flight-server
          chown flight:flight /opt/flight-server
          
          # Create systemd service file
          cat > /etc/systemd/system/flight-server.service << 'EOF'
          [Unit]
          Description=Apache Arrow Flight Server
          After=network.target
          
          [Service]
          Type=simple
          User=flight
          WorkingDirectory=/opt/flight-server
          ExecStart=/usr/bin/java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED -jar /opt/flight-server/simple-flight-server-1.0-SNAPSHOT.jar -server
          Restart=always
          RestartSec=10
          StandardOutput=journal
          StandardError=journal
          
          [Install]
          WantedBy=multi-user.target
          EOF
          
          # Enable the service (but don't start it yet - JAR needs to be uploaded first)
          systemctl enable flight-server
          
          # Install CloudWatch agent
          yum install -y amazon-cloudwatch-agent
          
          # Signal that the instance is ready
          /opt/aws/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource FlightServerInstance --region ${AWS::Region}
      Tags:
        - Key: Name
          Value: FlightServer-Instance
    CreationPolicy:
      ResourceSignal:
        Timeout: PT10M

  # Network Load Balancer
  NetworkLoadBalancer:
    Type: AWS::ElasticLoadBalancingV2::LoadBalancer
    Properties:
      Name: FlightServer-NLB
      Type: network
      Scheme: internet-facing
      Subnets:
        - !Ref PublicSubnet1
        - !Ref PublicSubnet2
      Tags:
        - Key: Name
          Value: FlightServer-NLB

  TargetGroup:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    Properties:
      Name: FlightServer-TG
      Port: 8815
      Protocol: TCP
      VpcId: !Ref VPC
      HealthCheckProtocol: TCP
      HealthCheckPort: 8815
      HealthCheckIntervalSeconds: 30
      HealthyThresholdCount: 2
      UnhealthyThresholdCount: 2
      Targets:
        - Id: !Ref FlightServerInstance
          Port: 8815
      Tags:
        - Key: Name
          Value: FlightServer-TG

  Listener:
    Type: AWS::ElasticLoadBalancingV2::Listener
    Properties:
      DefaultActions:
        - Type: forward
          TargetGroupArn: !Ref TargetGroup
      LoadBalancerArn: !Ref NetworkLoadBalancer
      Port: 8815
      Protocol: TCP
      ListenerAttributes:
        - Key: tcp.idle_timeout.seconds
          Value: '60'

Outputs:
  LoadBalancerDNS:
    Description: DNS name of the Network Load Balancer
    Value: !GetAtt NetworkLoadBalancer.DNSName
    Export:
      Name: !Sub "${AWS::StackName}-LoadBalancerDNS"
  
  LoadBalancerEndpoint:
    Description: Flight server endpoint (use this in your client)
    Value: !Sub "${NetworkLoadBalancer.DNSName}:8815"
    Export:
      Name: !Sub "${AWS::StackName}-FlightEndpoint"
  
  InstanceId:
    Description: EC2 Instance ID
    Value: !Ref FlightServerInstance
    Export:
      Name: !Sub "${AWS::StackName}-InstanceId"
  
  InstancePublicIP:
    Description: Public IP address of the EC2 instance (for SSH access)
    Value: !GetAtt FlightServerInstance.PublicIp
    Export:
      Name: !Sub "${AWS::StackName}-InstancePublicIP"
  
  SSHCommand:
    Description: SSH command to connect to the instance
    Value: !Sub "ssh -i your-key.pem ec2-user@${FlightServerInstance.PublicIp}"