ossService/oss/oss.go

package oss

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gctx"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"sync"
	"sync/atomic"

	"app.yhyue.com/moapp/jybase/es"
	ossSDK "github.com/aliyun/aliyun-oss-go-sdk/oss"
	"jygit.jydev.jianyu360.cn/BaseService/ossService/config"
	"jygit.jydev.jianyu360.cn/BaseService/ossService/util"
)

var (
	// accountMap 用来缓存OSS帐号信息
	accountMap           sync.Map
	DownLoadPool         = make(chan bool, g.Config().MustGet(gctx.New(), "downLoadPoolSize").Int())
	GetDetailFromEsPool  = make(chan bool, g.Config().MustGet(gctx.New(), "getDetailFromEsPoolSize").Int())
	GetDetailFromMgoPool = make(chan bool, g.Config().MustGet(gctx.New(), "getDetailFromMgoPoolSize").Int())
	getDetail            = map[string]func(bucketID, objectName string) []byte{
		"oss": func(bucketID, objectName string) []byte {
			b, _, e := DownloadAttachment(bucketID, objectName)
			if e == nil && b != nil && len(b) > 0 {
				return b
			}
			return nil
		},
		"es": func(bucketID, objectName string) []byte {
			GetDetailFromEsPool <- true
			defer func() {
				<-GetDetailFromEsPool
			}()
			indexName := g.Config().MustGet(gctx.New(), "elasticSearch.indexName").String()
			list := es.VarEs.Get(indexName, indexName, fmt.Sprintf(`{"query":{"bool":{"filter":{"term":{"_id":"%s"}}}},"_source":["detail"]}`, objectName))
			if list != nil && len(*list) > 0 {
				detail, _ := (*list)[0]["detail"].(string)
				return []byte(detail)
			}
			return nil
		},
		"mgo": func(bucketID, objectName string) []byte {
			GetDetailFromMgoPool <- true
			defer func() {
				<-GetDetailFromMgoPool
			}()
			data, _ := config.Mgo.FindOneByField(g.Config().MustGet(gctx.New(), "mongodb.collection").String(), objectName, `{"detail":1}`)
			if data != nil && len(*data) > 0 {
				detail, _ := (*data)["detail"].(string)
				return []byte(detail)
			}
			return nil
		},
	}
)

// GetBucket 根据bucketID获取bucket信息，如果没有则模拟查询数据库（这里只查询一次）
func GetBucket(bucketID string) (config.BucketInfo, error) {
	var buckets []*config.BucketInfo
	g.Config().MustGet(gctx.New(), "buckets").Structs(&buckets)
	for _, v := range buckets {
		if v.BucketID == bucketID {
			return *v, nil
		}
	}
	log.Println("Bucket %s not found in cache, querying database", bucketID)
	// 模拟数据库查询：实际业务中这里应查询数据库
	return config.BucketInfo{}, fmt.Errorf("bucket id %s not found", bucketID)
}

// bucketPool 缓存 bucket 对象，避免重复创建
var bucketPool sync.Map // key: bucketID, value: *ossSDK.Bucket

// GetCachedBucket 对外暴露统一入口，根据 bucketID 获取 *ossSDK.Bucket 对象
func GetCachedBucket(bucketID string) (*ossSDK.Bucket, error) {
	// 优先从对象池取
	if cached, ok := bucketPool.Load(bucketID); ok {
		return cached.(*ossSDK.Bucket), nil
	}

	// 如果不存在则依次读取 bucketInfo 和 OSS client 信息
	bucketInfo, err := GetBucket(bucketID)
	if err != nil {
		return nil, err
	}

	client, err := getOSSClient(bucketInfo.AccountID)
	if err != nil {
		return nil, err
	}

	bucket, err := client.Bucket(bucketInfo.BucketName)
	if err != nil {
		return nil, err
	}

	// 存入对象池
	bucketPool.Store(bucketID, bucket)
	return bucket, nil
}

// getOSSClient 根据OSS帐号ID从accountMap中获取OSS Client
func getOSSClient(accountID string) (*ossSDK.Client, error) {
	if val, ok := accountMap.Load(accountID); ok {
		acc := val.(config.OSSAccount)
		client, err := ossSDK.New(acc.Endpoint, acc.AccessKeyId, acc.AccessKeySecret)
		if err != nil {
			return nil, err
		}
		return client, nil
	}
	return nil, fmt.Errorf("oss account %s not found", accountID)
}

// UploadAttachment 上传附件；如果gzipEnabled为true，则进行gzip压缩（使用 io.Pipe 实现流式压缩）
func UploadAttachment(bucketID, objectName string, data io.Reader, gzipEnabled bool) error {
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return err
	}

	var reader io.Reader
	if gzipEnabled {
		pr, pw := io.Pipe()
		go func() {
			gzipWriter := gzip.NewWriter(pw)
			_, err := io.Copy(gzipWriter, data)
			// 注意：关闭 gzipWriter 时会将数据刷新到 pw 中
			gzipWriter.Close()
			pw.CloseWithError(err)
		}()
		reader = pr
	} else {
		reader = data
	}

	// 上传到OSS
	return bucket.PutObject(objectName, reader)
}

// DownloadAttachment 下载附件，如果检测到数据为gzip压缩，则自动解压后返回
func DownloadAttachment(bucketID, objectName string) ([]byte, http.Header, error) {
	atomic.AddInt64(&util.DownLoadCounter, 1)
	DownLoadPool <- true
	defer func() {
		<-DownLoadPool
		atomic.AddInt64(&util.DownLoadCounter, -1)
	}()
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return nil, nil, err
	}

	result, err := bucket.DoGetObject(&ossSDK.GetObjectRequest{objectName}, nil)
	if err != nil {
		return nil, nil, err
	}
	defer result.Response.Body.Close()
	data, err := ioutil.ReadAll(result.Response.Body)
	if err != nil {
		return nil, nil, err
	}
	// 判断是否为gzip压缩格式（判断前两个字节）
	if len(data) >= 2 && data[0] == 0x1f && data[1] == 0x8b {
		gzipReader, err := gzip.NewReader(bytes.NewReader(data))
		if err != nil {
			return nil, nil, err
		}
		defer gzipReader.Close()
		data, err = ioutil.ReadAll(gzipReader)
		if err != nil {
			return nil, nil, err
		}
		return data, result.Response.Headers, nil
	}
	return data, result.Response.Headers, nil
}

// DeleteAttachment 删除附件
func DeleteAttachment(bucketID, objectName string) error {
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return err
	}

	return bucket.DeleteObject(objectName)
}

// LoadOSSAccounts 加载OSS帐号信息到缓存
func LoadOSSAccounts() {
	var oas []*config.OSSAccount
	g.Config().MustGet(gctx.New(), "oss_accounts").Structs(&oas)
	accountMap.Clear()
	for _, acc := range oas {
		accountMap.Store(acc.ID, *acc)
	}
	bucketPool.Clear()
}

// 获取标讯正文，优先从oss中取，再从es中取
func GetBidDetail(bucketID, objectName string) []byte {
	atomic.AddInt64(&util.GetDetailCounter, 1)
	defer atomic.AddInt64(&util.GetDetailCounter, -1)
	for _, v := range g.Config().MustGet(gctx.New(), "getDetailOrder").Strings() {
		detail := getDetail[v](bucketID, objectName)
		if detail != nil && len(detail) > 0 {
			return detail
		}
	}
	return nil
}