Golang微服务RDS IAM认证令牌刷新失效问题求助

问题根源分析 #

1. 实例引用未更新：主函数中初始化的业务Repository持有了旧的GORM DB实例指针，即使刷新逻辑创建了新DB实例并更新了传入的指针，Repository仍会继续使用旧实例，无法获取新令牌对应的连接。
2. 资源泄漏风险：每次刷新都创建新DB实例，但旧实例的连接池未被正确关闭，会导致数据库连接资源泄漏。
3. AWS Session未复用：ConnectDataBase函数每次调用都新建AWS Session，没有利用AWS SDK自动刷新凭证的能力，增加了不必要的开销。

解决方案 #

1. 将DB实例托管到Service内部，提供线程安全的获取方式 #

通过Service持有DB实例，配合读写锁保证多线程下的实例替换与获取安全，让业务组件每次操作数据库时都能拿到最新的DB实例。

2. 优化令牌刷新逻辑 #

刷新时先创建新的DB实例，验证可用后再替换旧实例，并关闭旧连接池，避免资源泄漏；同时调整连接池生命周期参数，确保旧连接在令牌过期前被回收。

3. 复用AWS Session #

全局复用AWS Session，利用SDK自动处理凭证过期，无需手动重建Session。

代码修改示例 #

DB模块修改 #

package db

import (
    "fmt"
    "log"
    "sync"
    "time"

    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/rds/rdsutils"
    "github.com/go-co-op/gocron"
    "github.com/golang-migrate/migrate/v4"
    "github.com/golang-migrate/migrate/v4/database/postgres"
    _ "github.com/golang-migrate/migrate/v4/source/file"
    "github.com/jinzhu/gorm"
    _ "github.com/lib/pq"
    "your-project-path/common"
)

const DatabaseVersion = 1
const Success = "Database connection successful"
const database = "postgres"

// Service ...
type Service interface {
    InitDatabaseConnection()
    CreateDatabaseConnection()
    GetDB() *gorm.DB
    Migrate() error
}

type service struct {
    db *gorm.DB
    mu sync.RWMutex // 线程安全锁
}

var awsSession *session.Session

func init() {
    var err error
    awsSession, err = session.NewSession()
    if err != nil {
        panic(fmt.Sprintf("Failed to initialize AWS session: %v", err))
    }
}

// NewService ...
func NewService() Service {
    return &service{}
}

func (s *service) InitDatabaseConnection() {
    log.Print("Creating initial Db Connection")
    s.CreateDatabaseConnection()
    
    // 提前1分钟刷新（9分钟周期），避免令牌过期
    sch := gocron.NewScheduler(time.UTC)
    sch.Every(540).Seconds().Do(func() {
        log.Printf("Refreshing RDS auth credentials")
        s.CreateDatabaseConnection()
    })
    sch.StartAsync()
}

func (s *service) CreateDatabaseConnection() {
    configuration, err := common.New()
    if err != nil {
        log.Fatalf("Failed to load configuration: %v", err)
    }
    dbURI := ConnectDataBase(configuration)
    
    // 创建新DB实例
    newDB, err := gorm.Open(database, dbURI)
    if err != nil {
        log.Printf("Failed to create new DB connection: %v", err)
        return // 保留旧实例继续工作，避免服务崩溃
    }

    // 配置连接池参数，生命周期短于刷新周期
    newDB.DB().SetMaxOpenConns(6)
    newDB.DB().SetConnMaxLifetime(500 * time.Second)
    newDB.DB().SetMaxIdleConns(2)
    newDB.DB().SetConnMaxIdleTime(100 * time.Second)

    // 线程安全替换实例
    s.mu.Lock()
    defer s.mu.Unlock()

    // 关闭旧连接池
    if s.db != nil {
        if closeErr := s.db.Close(); closeErr != nil {
            log.Printf("Failed to close old DB connection pool: %v", closeErr)
        }
    }

    s.db = newDB
    log.Println(Success)
}

func (s *service) GetDB() *gorm.DB {
    s.mu.RLock()
    defer s.mu.RUnlock()
    return s.db
}

func ConnectDataBase(config *common.Constants) string {
    searchPath := config.SearchPath
    if searchPath == "" {
        searchPath = "public"
    }
    dbEndpoint := fmt.Sprintf("%s:%d", config.Host, config.Port)
    
    creds := awsSession.Config.Credentials
    authToken, err := rdsutils.BuildAuthToken(dbEndpoint, config.AwsRegion, config.UserName, creds)
    if err != nil {
        panic(fmt.Sprintf("Failed to build RDS auth token: %v", err))
    }

    return fmt.Sprintf("host=%s port=%d user=%s dbname=%s sslmode=require password=%s search_path=%s",
        config.Host, config.Port, config.UserName, config.DbName, authToken, searchPath)
}

func (s *service) Migrate() error {
    db := s.GetDB()
    if db == nil {
        return fmt.Errorf("database connection not initialized")
    }
    
    migrationSourceURL := "file://../resources/db/migrations/"
    database := db.DB()

    row := db.Table("schema_migrations").Limit(1).Row()
    var version int8
    var dirty bool

    err := row.Scan(&version, &dirty)
    if err == nil {
        log.Printf("Database is currently on version: %v", version)
    }

    log.Println("Migrating database to version", DatabaseVersion)
    driver, err := postgres.WithInstance(database, &postgres.Config{})
    if err != nil {
        return err
    }

    m, err := migrate.NewWithDatabaseInstance(migrationSourceURL, "postgres", driver)
    if err != nil {
        return err
    }

    m.Log = LogService{}
    err = m.Migrate(uint(DatabaseVersion))
    if err != nil && err.Error() != "no change" {
        log.Println("Migration error:", err.Error())
        return err
    }
    return nil
}

Main模块修改 #

package main

import (
    "bufio"
    "encoding/base64"
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "os"
    "strings"
    "time"

    gokitlog "github.com/go-kit/kit/log"
    "github.com/go-kit/kit/log/level"
    "github.com/gorilla/mux"
    "your-project-path/db"
    "your-project-path/common"
    "your-project-path/allocation"
)

func runServer(config *common.Constants) {
    var logger gokitlog.Logger
    logger = gokitlog.NewLogfmtLogger(gokitlog.NewSyncWriter(os.Stderr))
    logger = gokitlog.With(logger, "ts", gokitlog.DefaultTimestampUTC)

    dbService := db.NewService()
    dbService.InitDatabaseConnection()

    err := dbService.Migrate()
    if err != nil {
        fmt.Println("Migration scripts failed to run:", err.Error())
        os.Exit(1)
    }
    level.Info(logger).Log("Starting Darwin ==========>")

    // 传递DB Service给Repository，而非直接传递DB实例
    repo, err := allocation.NewUserAllocationRepository(dbService, logger)
    if err != nil {
        level.Error(logger).Log("exit", err)
        os.Exit(3)
    }
    userAllocationService := allocation.NewUserAllocationService(repo, logger)

    endpoints := allocation.MakeEndpoints(userAllocationService)

    router := mux.NewRouter()
    sh := http.StripPrefix("/swaggerui/", http.FileServer(http.Dir("../resources/swaggerui/")))
    router.PathPrefix("/swaggerui/").Handler(sh)
    httpLogger := gokitlog.With(logger, "component", "http")
    subRouter := router.PathPrefix("/api/v1").Subrouter()
    subRouter.NotFoundHandler = http.HandlerFunc(notFound)
    subRouter.Use(checkBasicAuth(config))
    subRouter.HandleFunc("/healthCheck", healthCheck).Methods("GET")
    allocation.MakeHandler(subRouter, httpLogger, endpoints)
    http.Handle("/", subRouter)

    log.Print("Running server on port 8080")
    f, err := os.OpenFile("../performance.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
    if err != nil {
        log.Printf("Error opening performance log: %v", err)
        log.Fatal(http.ListenAndServe(":8080", router))
    } else {
        by := bufio.NewWriter(f)
        defer f.Close()
        log.Fatal(http.ListenAndServe(":8080", Logger(by, router)))
    }
}

// 其余函数保持不变...

Repository修改示例 #

package allocation

import (
    gokitlog "github.com/go-kit/kit/log"
    "github.com/jinzhu/gorm"
    "your-project-path/db"
)

type UserAllocationRepository struct {
    dbService db.Service
    logger    gokitlog.Logger
}

func NewUserAllocationRepository(dbService db.Service, logger gokitlog.Logger) (*UserAllocationRepository, error) {
    return &UserAllocationRepository{
        dbService: dbService,
        logger:    logger,
    }, nil
}

// 每次操作DB时获取最新实例
func (r *UserAllocationRepository) GetUserAllocation(id int) (*Allocation, error) {
    db := r.dbService.GetDB()
    if db == nil {
        return nil, fmt.Errorf("database connection unavailable")
    }
    
    var allocation Allocation
    err := db.Where("id = ?", id).First(&allocation).Error
    return &allocation, err
}

内容的提问来源于stack exchange，提问作者Abhisek Roy