diff --git a/integration_tests/cdk/app.py b/integration_tests/cdk/app.py index 1a04621..8020c27 100644 --- a/integration_tests/cdk/app.py +++ b/integration_tests/cdk/app.py @@ -79,6 +79,7 @@ def __init__( ), allocated_storage=app_config.db_allocated_storage, instance_type=aws_ec2.InstanceType(app_config.db_instance_type), + add_pgbouncer=True, removal_policy=RemovalPolicy.DESTROY, ) diff --git a/lib/database/PgBouncer.ts b/lib/database/PgBouncer.ts new file mode 100644 index 0000000..c06a87a --- /dev/null +++ b/lib/database/PgBouncer.ts @@ -0,0 +1,270 @@ +import { + aws_ec2 as ec2, + aws_iam as iam, + aws_lambda as lambda, + aws_secretsmanager as secretsmanager, + CustomResource, + Stack, +} from "aws-cdk-lib"; +import { Construct } from "constructs"; + +import * as fs from "fs"; +import * as path from "path"; + +// used to populate pgbouncer config: +// see https://www.pgbouncer.org/config.html for details +export interface PgBouncerConfigProps { + poolMode?: "transaction" | "session" | "statement"; + maxClientConn?: number; + defaultPoolSize?: number; + minPoolSize?: number; + reservePoolSize?: number; + reservePoolTimeout?: number; + maxDbConnections?: number; + maxUserConnections?: number; +} + +export interface PgBouncerProps { + /** + * Name for the pgbouncer instance + */ + instanceName: string; + + /** + * VPC to deploy PgBouncer into + */ + vpc: ec2.IVpc; + + /** + * The RDS instance to connect to + */ + database: { + connections: ec2.Connections; + secret: secretsmanager.ISecret; + }; + + /** + * Maximum connections setting for the database + */ + dbMaxConnections: number; + + /** + * Whether to deploy in public subnet + * @default false + */ + usePublicSubnet?: boolean; + + /** + * Instance type for PgBouncer + * @default t3.micro + */ + instanceType?: ec2.InstanceType; + + /** + * PgBouncer configuration options + */ + pgBouncerConfig?: PgBouncerConfigProps; +} + +export class PgBouncer extends Construct { + public readonly instance: ec2.Instance; + public readonly pgbouncerSecret: secretsmanager.Secret; + + // The max_connections parameter in PgBouncer determines the maximum number of + // connections to open on the actual database instance. We want that number to + // be slightly smaller than the actual max_connections value on the RDS instance + // so we perform this calculation. + + private getDefaultConfig( + dbMaxConnections: number + ): Required { + // maxDbConnections (and maxUserConnections) are the only settings that need + // to be responsive to the database size/max_connections setting + return { + poolMode: "transaction", + maxClientConn: 1000, + defaultPoolSize: 5, + minPoolSize: 0, + reservePoolSize: 5, + reservePoolTimeout: 5, + maxDbConnections: dbMaxConnections - 10, + maxUserConnections: dbMaxConnections - 10, + }; + } + + constructor(scope: Construct, id: string, props: PgBouncerProps) { + super(scope, id); + + // Set defaults for optional props + const defaultInstanceType = ec2.InstanceType.of( + ec2.InstanceClass.T3, + ec2.InstanceSize.MICRO + ); + + const instanceType = props.instanceType ?? defaultInstanceType; + const defaultConfig = this.getDefaultConfig(props.dbMaxConnections); + + // Merge provided config with defaults + const pgBouncerConfig: Required = { + ...defaultConfig, + ...props.pgBouncerConfig, + }; + + // Create role for PgBouncer instance to enable writing to CloudWatch + const role = new iam.Role(this, "InstanceRole", { + description: + "pgbouncer instance role with Systems Manager + CloudWatch permissions", + assumedBy: new iam.ServicePrincipal("ec2.amazonaws.com"), + managedPolicies: [ + iam.ManagedPolicy.fromAwsManagedPolicyName( + "AmazonSSMManagedInstanceCore" + ), + iam.ManagedPolicy.fromAwsManagedPolicyName( + "CloudWatchAgentServerPolicy" + ), + ], + }); + + // Add policy to allow reading RDS credentials from Secrets Manager + role.addToPolicy( + new iam.PolicyStatement({ + actions: ["secretsmanager:GetSecretValue"], + resources: [props.database.secret.secretArn], + }) + ); + + // Create PgBouncer instance + this.instance = new ec2.Instance(this, "Instance", { + vpc: props.vpc, + vpcSubnets: { + subnetType: props.usePublicSubnet + ? ec2.SubnetType.PUBLIC + : ec2.SubnetType.PRIVATE_WITH_EGRESS, + }, + instanceType, + instanceName: props.instanceName, + machineImage: ec2.MachineImage.fromSsmParameter( + "/aws/service/canonical/ubuntu/server/jammy/stable/current/amd64/hvm/ebs-gp2/ami-id", + { os: ec2.OperatingSystemType.LINUX } + ), + role, + blockDevices: [ + { + deviceName: "/dev/xvda", + volume: ec2.BlockDeviceVolume.ebs(20, { + volumeType: ec2.EbsDeviceVolumeType.GP3, + encrypted: true, + deleteOnTermination: true, + }), + }, + ], + userData: this.loadUserDataScript(pgBouncerConfig, props.database), + userDataCausesReplacement: true, + }); + + // Allow PgBouncer to connect to RDS + props.database.connections.allowFrom( + this.instance, + ec2.Port.tcp(5432), + "Allow PgBouncer to connect to RDS" + ); + + // Create a new secret for pgbouncer connection credentials + this.pgbouncerSecret = new secretsmanager.Secret(this, "PgBouncerSecret", { + description: `Connection information for PgBouncer instance ${props.instanceName}`, + generateSecretString: { + generateStringKey: "dummy", + secretStringTemplate: "{}", + }, + }); + + // Grant the role permission to read the new secret + this.pgbouncerSecret.grantRead(role); + + // Update pgbouncerSecret to contain pgstacSecret values but with new value for host + const secretUpdaterFn = new lambda.Function(this, "SecretUpdaterFunction", { + runtime: lambda.Runtime.NODEJS_20_X, + handler: "index.handler", + code: lambda.Code.fromInline(` + const AWS = require('aws-sdk'); + const sm = new AWS.SecretsManager(); + + exports.handler = async (event) => { + console.log('Event:', JSON.stringify(event, null, 2)); + + try { + const instanceIp = event.ResourceProperties.instanceIp; + + // Get the original secret value + const originalSecret = await sm.getSecretValue({ + SecretId: '${props.database.secret.secretArn}' + }).promise(); + + // Parse the secret string + const secretData = JSON.parse(originalSecret.SecretString); + + // Update the host value with the PgBouncer instance IP + secretData.host = instanceIp; + + // Put the modified secret value + await sm.putSecretValue({ + SecretId: '${this.pgbouncerSecret.secretArn}', + SecretString: JSON.stringify(secretData) + }).promise(); + + return { + PhysicalResourceId: '${this.pgbouncerSecret.secretArn}', + Data: { + SecretArn: '${this.pgbouncerSecret.secretArn}' + } + }; + } catch (error) { + console.error('Error:', error); + throw error; + } + }; + `), + }); + + props.database.secret.grantRead(secretUpdaterFn); + this.pgbouncerSecret.grantWrite(secretUpdaterFn); + + new CustomResource(this, "pgbouncerSecretBootstrapper", { + serviceToken: secretUpdaterFn.functionArn, + properties: { + instanceIp: this.instance.instancePrivateIp, + }, + }); + } + + private loadUserDataScript( + pgBouncerConfig: Required>, + database: { secret: secretsmanager.ISecret } + ): ec2.UserData { + const userDataScript = ec2.UserData.forLinux(); + + // Set environment variables with configuration parameters + userDataScript.addCommands( + 'export SECRET_ARN="' + database.secret.secretArn + '"', + 'export REGION="' + Stack.of(this).region + '"', + 'export POOL_MODE="' + pgBouncerConfig.poolMode + '"', + 'export MAX_CLIENT_CONN="' + pgBouncerConfig.maxClientConn + '"', + 'export DEFAULT_POOL_SIZE="' + pgBouncerConfig.defaultPoolSize + '"', + 'export MIN_POOL_SIZE="' + pgBouncerConfig.minPoolSize + '"', + 'export RESERVE_POOL_SIZE="' + pgBouncerConfig.reservePoolSize + '"', + 'export RESERVE_POOL_TIMEOUT="' + + pgBouncerConfig.reservePoolTimeout + + '"', + 'export MAX_DB_CONNECTIONS="' + pgBouncerConfig.maxDbConnections + '"', + 'export MAX_USER_CONNECTIONS="' + pgBouncerConfig.maxUserConnections + '"' + ); + + // Load the startup script + const scriptPath = path.join(__dirname, "./pgbouncer-setup.sh"); + let script = fs.readFileSync(scriptPath, "utf8"); + + userDataScript.addCommands(script); + + return userDataScript; + } +} diff --git a/lib/database/index.ts b/lib/database/index.ts index 18d6444..f7e8caf 100644 --- a/lib/database/index.ts +++ b/lib/database/index.ts @@ -8,18 +8,18 @@ import { RemovalPolicy, Duration, aws_logs, - } from "aws-cdk-lib"; import { Construct } from "constructs"; import { CustomLambdaFunctionProps } from "../utils"; +import { PgBouncer } from "./PgBouncer"; const instanceSizes: Record = require("./instance-memory.json"); const DEFAULT_PGSTAC_VERSION = "0.8.5"; -let defaultPgSTACCustomOptions :{ [key: string]: any } = { - "context": "FALSE", - "mosaic_index": "TRUE" -} +let defaultPgSTACCustomOptions: { [key: string]: any } = { + context: "FALSE", + mosaic_index: "TRUE", +}; function hasVpc( instance: rds.DatabaseInstance | rds.IDatabaseInstance @@ -35,6 +35,7 @@ function hasVpc( export class PgStacDatabase extends Construct { db: rds.DatabaseInstance; pgstacSecret: secretsmanager.ISecret; + pgbouncerServer?: PgBouncer; constructor(scope: Construct, id: string, props: PgStacDatabaseProps) { super(scope, id); @@ -46,6 +47,7 @@ export class PgStacDatabase extends Construct { const parameterGroup = new rds.ParameterGroup(this, "parameterGroup", { engine: props.engine, parameters: { + max_connections: defaultParameters.maxConnections, shared_buffers: defaultParameters.sharedBuffers, effective_cache_size: defaultParameters.effectiveCacheSize, work_mem: defaultParameters.workMem, @@ -73,7 +75,10 @@ export class PgStacDatabase extends Construct { timeout: Duration.minutes(2), code: aws_lambda.Code.fromDockerBuild(__dirname, { file: "bootstrapper_runtime/Dockerfile", - buildArgs: {PGSTAC_VERSION: DEFAULT_PGSTAC_VERSION, PYTHON_VERSION: "3.11"} + buildArgs: { + PGSTAC_VERSION: DEFAULT_PGSTAC_VERSION, + PYTHON_VERSION: "3.11", + }, }), vpc: hasVpc(this.db) ? this.db.vpc : props.vpc, allowPublicSubnet: true, @@ -111,11 +116,15 @@ export class PgStacDatabase extends Construct { // connect to database this.db.connections.allowFrom(handler, ec2.Port.tcp(5432)); - let customResourceProperties : { [key: string]: any} = props.customResourceProperties ? { ...defaultPgSTACCustomOptions, ...props.customResourceProperties } : defaultPgSTACCustomOptions; + let customResourceProperties: { [key: string]: any } = + props.customResourceProperties + ? { ...defaultPgSTACCustomOptions, ...props.customResourceProperties } + : defaultPgSTACCustomOptions; // update properties customResourceProperties["conn_secret_arn"] = this.db.secret!.secretArn; - customResourceProperties["new_user_secret_arn"] = this.pgstacSecret.secretArn; + customResourceProperties["new_user_secret_arn"] = + this.pgstacSecret.secretArn; // if props.lambdaFunctionOptions doesn't have 'code' defined, update pgstac_version (needed for default runtime) if (!props.bootstrapperLambdaFunctionOptions?.code) { @@ -128,6 +137,33 @@ export class PgStacDatabase extends Construct { removalPolicy: RemovalPolicy.RETAIN, // This retains the custom resource (which doesn't really exist), not the database }); + // PgBouncer + const addPgbouncer = props.addPgbouncer ?? true; + if (addPgbouncer) { + this.pgbouncerServer = new PgBouncer(this, "pgbouncer", { + instanceName: `${Stack.of(this).stackName}-pgbouncer`, + instanceType: ec2.InstanceType.of( + ec2.InstanceClass.T3, + ec2.InstanceSize.MICRO + ), + vpc: props.vpc, + database: { + connections: this.db.connections, + secret: this.pgstacSecret, + }, + dbMaxConnections: parseInt(defaultParameters.maxConnections), + usePublicSubnet: false, + pgBouncerConfig: { + poolMode: "transaction", + maxClientConn: 1000, + defaultPoolSize: 20, + minPoolSize: 10, + reservePoolSize: 5, + reservePoolTimeout: 5, + }, + }); + this.pgstacSecret = this.pgbouncerServer.pgbouncerSecret; + } } public getParameters( @@ -178,12 +214,12 @@ export interface PgStacDatabaseProps extends rds.DatabaseInstanceProps { */ readonly pgstacDbName?: string; - /** + /** * Prefix to assign to the generated `secrets_manager.Secret` * * @default pgstac */ - readonly secretsPrefix?: string; + readonly secretsPrefix?: string; /** * Name of user that will be generated for connecting to the pgSTAC database. @@ -192,6 +228,13 @@ export interface PgStacDatabaseProps extends rds.DatabaseInstanceProps { */ readonly pgstacUsername?: string; + /** + * Add pgbouncer instance for managing traffic to the pgSTAC database + * + * @default true + */ + readonly addPgbouncer?: boolean; + /** * Lambda function Custom Resource properties. A custom resource property is going to be created * to trigger the boostrapping lambda function. This parameter allows the user to specify additional properties @@ -200,7 +243,7 @@ export interface PgStacDatabaseProps extends rds.DatabaseInstanceProps { */ readonly customResourceProperties?: { [key: string]: any; -} + }; /** * Can be used to override the default lambda function properties. diff --git a/lib/database/pgbouncer-setup.sh b/lib/database/pgbouncer-setup.sh new file mode 100644 index 0000000..3aa6d01 --- /dev/null +++ b/lib/database/pgbouncer-setup.sh @@ -0,0 +1,257 @@ +#!/bin/bash +set -euxo pipefail + +# These variables will be replaced by the TypeScript code +SECRET_ARN=${SECRET_ARN} +REGION=${REGION} +POOL_MODE=${POOL_MODE} +MAX_CLIENT_CONN=${MAX_CLIENT_CONN} +DEFAULT_POOL_SIZE=${DEFAULT_POOL_SIZE} +MIN_POOL_SIZE=${MIN_POOL_SIZE} +RESERVE_POOL_SIZE=${RESERVE_POOL_SIZE} +RESERVE_POOL_TIMEOUT=${RESERVE_POOL_TIMEOUT} +MAX_DB_CONNECTIONS=${MAX_DB_CONNECTIONS} +MAX_USER_CONNECTIONS=${MAX_USER_CONNECTIONS} +CLOUDWATCH_CONFIG="/opt/aws/amazon-cloudwatch-agent/bin/config.json" + +# Add the postgres repository +curl https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add - +sudo sh -c 'echo "deb http://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list' + +# Install required packages +apt-get update + +# Function that makes sure we don't hit a dpkg lock error +wait_for_dpkg_lock() { + while fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock >/dev/null 2>&1; do + echo "Waiting for dpkg lock to be released..." + sleep 2 + done +} + +wait_for_dpkg_lock +DEBIAN_FRONTEND=noninteractive apt-get install -y pgbouncer jq awscli + +echo "Fetching secret from ARN: ${SECRET_ARN}" + +# Before handling secrets, turn off command tracing +set +x +SECRET=$(aws secretsmanager get-secret-value --secret-id ${SECRET_ARN} --region ${REGION} --query SecretString --output text) + +# Parse database credentials without echoing +DB_HOST=$(echo "$SECRET" | jq -r '.host') +DB_PORT=$(echo "$SECRET" | jq -r '.port') +DB_NAME=$(echo "$SECRET" | jq -r '.dbname') +DB_USER=$(echo "$SECRET" | jq -r '.username') +DB_PASSWORD=$(echo "$SECRET" | jq -r '.password') + +echo 'Creating PgBouncer configuration...' + +# Create pgbouncer.ini +cat < /etc/pgbouncer/pgbouncer.ini +[databases] +* = host=$DB_HOST port=$DB_PORT dbname=$DB_NAME + +[pgbouncer] +listen_addr = 0.0.0.0 +listen_port = 5432 +auth_type = md5 +auth_file = /etc/pgbouncer/userlist.txt +pool_mode = ${POOL_MODE} +max_client_conn = ${MAX_CLIENT_CONN} +default_pool_size = ${DEFAULT_POOL_SIZE} +min_pool_size = ${MIN_POOL_SIZE} +reserve_pool_size = ${RESERVE_POOL_SIZE} +reserve_pool_timeout = ${RESERVE_POOL_TIMEOUT} +max_db_connections = ${MAX_DB_CONNECTIONS} +max_user_connections = ${MAX_USER_CONNECTIONS} +max_prepared_statements = 10 +ignore_startup_parameters = application_name,search_path +logfile = /var/log/pgbouncer/pgbouncer.log +pidfile = /var/run/pgbouncer/pgbouncer.pid +admin_users = $DB_USER +stats_users = $DB_USER +log_connections = 1 +log_disconnections = 1 +log_pooler_errors = 1 +log_stats = 1 +stats_period = 60 +EOC + +# Create userlist.txt without echoing sensitive info +{ + echo "\"$DB_USER\" \"$DB_PASSWORD\"" +} > /etc/pgbouncer/userlist.txt + +# Turn command tracing back on +set -x + +# Set correct permissions +chown postgres:postgres /etc/pgbouncer/pgbouncer.ini /etc/pgbouncer/userlist.txt +chmod 600 /etc/pgbouncer/pgbouncer.ini /etc/pgbouncer/userlist.txt + +# Configure logging +# ensure /var/run/pgbouncer gets created on boot +cat < /etc/tmpfiles.d/pgbouncer.conf +d /var/run/pgbouncer 0755 postgres postgres - +EOC + +mkdir -p /var/log/pgbouncer /var/run/pgbouncer +chown postgres:postgres /var/log/pgbouncer /var/run/pgbouncer +chmod 755 /var/log/pgbouncer /var/run/pgbouncer + +touch /var/log/pgbouncer/pgbouncer.log +chown postgres:postgres /var/log/pgbouncer/pgbouncer.log +chmod 640 /var/log/pgbouncer/pgbouncer.log + +# Enable and start pgbouncer service +systemctl enable pgbouncer +systemctl restart pgbouncer + + +cat < /etc/logrotate.d/pgbouncer +/var/log/pgbouncer/pgbouncer.log { + daily + rotate 7 + compress + delaycompress + missingok + copytruncate + create 640 postgres postgres +} +EOC + +# Create monitoring scripts directory +mkdir -p /opt/pgbouncer/scripts + +# Create the health check script +cat <<'EOC' > /opt/pgbouncer/scripts/check.sh +#!/bin/bash +echo $(/bin/systemctl is-active pgbouncer) +if ! /bin/systemctl is-active --quiet pgbouncer; then + # If it's not active, attempt to start it + echo "$(date): PgBouncer is not running, attempting to restart" | logger -t pgbouncer-monitor + /bin/systemctl start pgbouncer + + # Check if the restart was successful + if /bin/systemctl is-active --quiet pgbouncer; then + echo "$(date): PgBouncer successfully restarted" | logger -t pgbouncer-monitor + else + echo "$(date): Failed to restart PgBouncer" | logger -t pgbouncer-monitor + fi +else + # If it's already active, no action is needed + echo "$(date): PgBouncer is running; no action needed" | logger -t pgbouncer-monitor +fi +EOC +chmod +x /opt/pgbouncer/scripts/check.sh + +# enable cron job +cat <<'EOC' > /opt/pgbouncer/scripts/crontab.txt +# PgBouncer health check - run every minute +* * * * * /opt/pgbouncer/scripts/check.sh +EOC + +crontab /opt/pgbouncer/scripts/crontab.txt + +if ! crontab -l; then + echo 'Failed to install crontab' | logger -t pgbouncer-setup + exit 1 +fi + +# Create CloudWatch configuration directory +mkdir -p /opt/pgbouncer/cloudwatch + +# Install CloudWatch agent +if ! wget -q https://s3.amazonaws.com/amazoncloudwatch-agent/ubuntu/amd64/latest/amazon-cloudwatch-agent.deb; then + echo 'Failed to download CloudWatch agent' | logger -t pgbouncer-setup + exit 1 +fi + +if ! dpkg -i amazon-cloudwatch-agent.deb; then + echo 'Failed to install CloudWatch agent' | logger -t pgbouncer-setup + exit 1 +fi + +# Create CloudWatch config +cat < ${CLOUDWATCH_CONFIG} +{ + "agent": { + "metrics_collection_interval": 60, + "run_as_user": "root" + }, + "logs": { + "logs_collected": { + "files": { + "collect_list": [ + { + "file_path": "/var/log/pgbouncer/pgbouncer.log", + "log_group_name": "/pgbouncer/logs", + "log_stream_name": "{instance_id}", + "timestamp_format": "%Y-%m-%d %H:%M:%S", + "multi_line_start_pattern": "{timestamp_format}", + "retention_in_days": 14 + }, + { + "file_path": "/var/log/syslog", + "log_group_name": "/pgbouncer/system-logs", + "log_stream_name": "{instance_id}", + "timestamp_format": "%b %d %H:%M:%S", + "retention_in_days": 14 + } + ] + } + } + }, + "metrics": { + "metrics_collected": { + "procstat": [ + { + "pattern": "pgbouncer", + "measurement": [ + "cpu_usage", + "memory_rss", + "read_bytes", + "write_bytes", + "read_count", + "write_count", + "num_fds" + ] + } + ], + "mem": { + "measurement": [ + "mem_used_percent" + ] + }, + "disk": { + "measurement": [ + "used_percent" + ] + } + }, + "aggregation_dimensions": [["InstanceId"]] + } +} +EOC + +# Verify the config file exists +if [ ! -f ${CLOUDWATCH_CONFIG} ]; then + echo 'CloudWatch config file not created' | logger -t pgbouncer-setup + exit 1 +fi + +# Start CloudWatch agent +if ! /opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl -a fetch-config -m ec2 -s -c file:${CLOUDWATCH_CONFIG}; then + echo 'Failed to configure CloudWatch agent' | logger -t pgbouncer-setup + exit 1 +fi + +systemctl enable amazon-cloudwatch-agent +systemctl start amazon-cloudwatch-agent + +# Verify CloudWatch agent is running +if ! systemctl is-active amazon-cloudwatch-agent; then + echo 'CloudWatch agent failed to start' | logger -t pgbouncer-setup + exit 1 +fi