VptPassiveAdapter/fs/s3_adapter.go

package fs

import (
	"ZhenTuLocalPassiveAdapter/config"
	"ZhenTuLocalPassiveAdapter/dto"
	"ZhenTuLocalPassiveAdapter/util"
	"context"
	"fmt"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"log"
	"path"
	"sort"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
)

var s3Cache sync.Map

type S3Adapter struct {
	StorageConfig config.StorageConfig
	s3Client      *s3.Client
}

func (s *S3Adapter) getClient() (*s3.Client, error) {
	if s.s3Client == nil {
		const defaultRegion = "us-east-1"
		resolver := aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) {
			return aws.Endpoint{
				PartitionID:       "aws",
				URL:               s.StorageConfig.S3.Endpoint, // or where ever you ran minio
				SigningRegion:     defaultRegion,
				HostnameImmutable: true,
			}, nil
		})
		creds := credentials.NewStaticCredentialsProvider(s.StorageConfig.S3.AkId, s.StorageConfig.S3.AkSec, "")
		cfg := aws.Config{
			Credentials:      creds,
			Region:           defaultRegion,
			EndpointResolver: resolver,
		}
		s.s3Client = s3.NewFromConfig(cfg)
	}
	return s.s3Client, nil
}

func (s *S3Adapter) GetFileList(ctx context.Context, dirPath string, relDt time.Time) ([]dto.File, error) {
	_, span := tracer.Start(ctx, "GetFileList_s3")
	defer span.End()

	span.SetAttributes(attribute.String("path", dirPath))
	span.SetAttributes(attribute.String("relativeDate", relDt.Format("2006-01-02")))
	if s.StorageConfig.S3.Bucket == "" {
		span.SetAttributes(attribute.String("error", "未配置S3存储桶"))
		span.SetStatus(codes.Error, "未配置S3存储桶")
		return nil, fmt.Errorf("未配置S3存储桶")
	}

	cacheKey := fmt.Sprintf("%s_%s", dirPath, relDt.Format("2006-01-02"))
	if cachedInterface, ok := s3Cache.Load(cacheKey); ok {
		cachedItem := cachedInterface.(cacheItem)
		log.Println("缓存过期时间", cachedItem.expires.Sub(time.Now()))
		if time.Now().Before(cachedItem.expires) {
			log.Println("获取已缓存列表", cacheKey)
			span.SetAttributes(attribute.Bool("cache.hit", true))
			return cachedItem.data, nil
		}
	}

	mutexKey := fmt.Sprintf("lock_%s", cacheKey)
	mutex, _ := s3Cache.LoadOrStore(mutexKey, &sync.Mutex{})
	lock := mutex.(*sync.Mutex)
	defer func() {
		// 解锁后删除锁（避免内存泄漏）
		s3Cache.Delete(mutexKey)
		lock.Unlock()
	}()
	lock.Lock()

	if cachedInterface, ok := s3Cache.Load(cacheKey); ok {
		cachedItem := cachedInterface.(cacheItem)
		log.Println("缓存过期时间", cachedItem.expires.Sub(time.Now()))
		if time.Now().Before(cachedItem.expires) {
			log.Println("过锁后获取已缓存列表", cacheKey)
			span.SetAttributes(attribute.Bool("s3Cache.hit", true))
			return cachedItem.data, nil
		}
	}

	listObjectsInput := &s3.ListObjectsV2Input{
		Bucket:  aws.String(s.StorageConfig.S3.Bucket),
		Prefix:  aws.String(path.Join(s.StorageConfig.S3.Prefix, dirPath)),
		MaxKeys: aws.Int32(1000),
	}

	client, err := s.getClient()
	if err != nil {
		span.SetAttributes(attribute.String("error", err.Error()))
		span.SetStatus(codes.Error, "创建S3客户端失败")
		return nil, err
	}

	var fileList []dto.File
	var continuationToken *string

	for {
		if continuationToken != nil {
			listObjectsInput.ContinuationToken = continuationToken
		}

		result, err := client.ListObjectsV2(context.TODO(), listObjectsInput)
		if err != nil {
			span.SetAttributes(attribute.String("error", err.Error()))
			span.SetStatus(codes.Error, "文件列表读取失败")
			return nil, err
		}

		for _, object := range result.Contents {
			key := *object.Key
			if !util.IsVideoFile(path.Base(key)) {
				continue
			}
			startTime, stopTime, err := util.ParseStartStopTime(path.Base(key), relDt)
			if err != nil {
				continue
			}
			if stopTime.IsZero() {
				stopTime = startTime
			}
			if startTime.Equal(stopTime) {
				stopTime = stopTime.Add(time.Second * time.Duration(config.Config.Record.Duration))
			}
			presignClient := s3.NewPresignClient(client)
			request, err := presignClient.PresignGetObject(context.TODO(), &s3.GetObjectInput{
				Bucket: aws.String(s.StorageConfig.S3.Bucket),
				Key:    aws.String(key),
			}, func(presignOptions *s3.PresignOptions) {
				presignOptions.Expires = 10 * time.Minute
			})
			if err != nil {
				span.SetAttributes(attribute.String("error", err.Error()))
				span.SetStatus(codes.Error, "生成预签名URL失败")
				log.Println("Error presigning GetObject request:", err)
				continue
			}
			fileList = append(fileList, dto.File{
				BasePath:  s.StorageConfig.S3.Bucket,
				Name:      path.Base(key),
				Path:      path.Dir(key),
				Url:       request.URL,
				StartTime: startTime,
				EndTime:   stopTime,
			})
		}

		if !*result.IsTruncated {
			break
		}
		continuationToken = result.NextContinuationToken
	}

	span.SetAttributes(attribute.Int("file.count", len(fileList)))
	sort.Slice(fileList, func(i, j int) bool {
		return fileList[i].StartTime.Before(fileList[j].StartTime)
	})
	span.SetStatus(codes.Ok, "文件读取成功")

	cacheItem := cacheItem{
		data:    fileList,
		expires: time.Now().Add(10 * time.Second),
	}
	s3Cache.Store(cacheKey, cacheItem)
	log.Println("缓存文件列表", cacheKey)

	return fileList, nil
}

type cacheItem struct {
	data    []dto.File
	expires time.Time
}

// 添加定时清理缓存的初始化函数
func init() {
	go func() {
		ticker := time.NewTicker(1 * time.Minute)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				cleanupCache()
			}
		}
	}()
}

// 添加缓存清理函数
func cleanupCache() {
	var keysToDelete []interface{}
	s3Cache.Range(func(key, value interface{}) bool {
		item := value.(cacheItem)
		if time.Now().After(item.expires) {
			keysToDelete = append(keysToDelete, key)
		}
		return true
	})
	for _, key := range keysToDelete {
		s3Cache.Delete(key)
	}
}