hadja_server/pub/common/linode/linode.go

package linode

import (
	"bytes"
	"errors"
	"io"
	"mime"
	"os"
	"path"
	"path/filepath"
	"videoweb/pub/common/fconst"
	"videoweb/pub/common/log"

	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/karrick/godirwalk"
	"github.com/xxjwxc/gowp/workpool"
)

const (
	//region             = "default"
	//endpoint           = "https://ap-south-1.linodeobjects.com"
	//accessKey          = "0RALF74H45ANVYPRQA1Y"
	//secretKey          = "S1yCKUqxFJHqlFPZXJ3BbD8T6P0Yz6loUDntUrsS"
	//bucket             = "fengshang"
	M                  = 1024 * 1024
	G                  = M * 1024
	defaultConcurrency = 150
	defaultTryMax      = 5
	defaultContentType = "application/octet-stream"
)

type Client struct {
	session *session.Session
	cfg     Config
}
type UploadReq struct {
	Key         string
	ContentType string
	Data        []byte
}

func (u *UploadReq) Valid() bool {
	return u.Key != "" && len(u.Data) != 0
}

type Config struct {
	Region      string
	Endpoint    string
	AccessKey   string
	SecretKey   string
	Bucket      string
	Concurrency int
	TryMax      int
}

func traverseDir(d string) ([]string, error) {
	var fileNames []string
	fn := func(name string, entry *godirwalk.Dirent) error {
		if !entry.IsRegular() {
			return nil
		}
		fileNames = append(fileNames, name)
		return nil
	}
	pathName := d
	err := godirwalk.Walk(pathName, &godirwalk.Options{
		Unsorted: true,
		Callback: fn,
	})
	if err != nil {
		log.Info("list files failed", log.Any("path", pathName), log.E(err))
		return nil, err
	}
	return fileNames, nil
}

func New(cfg Config) (*Client, error) {
	if cfg.TryMax <= 0 {
		cfg.TryMax = defaultTryMax
	}
	if cfg.Concurrency <= 0 {
		cfg.Concurrency = defaultConcurrency
	}
	s, err := session.NewSession(&aws.Config{
		Region:      aws.String(cfg.Region),
		Endpoint:    aws.String(cfg.Endpoint),
		Credentials: credentials.NewStaticCredentials(cfg.AccessKey, cfg.SecretKey, ""),
	})
	if err != nil {
		log.Error("Linode NewSession err", log.E(err))
		return nil, err
	}
	return &Client{session: s, cfg: cfg}, nil
}

//func WithDefault() (*Client, error) {
//	return New(&Config{Concurrency: defaultConcurrency, TryMax: defaultTryMax})
//}

func (c *Client) UploadFolder(folderName string) error {
	absFiles, err := traverseDir(folderName)
	if err != nil {
		return err
	}
	if len(absFiles) == 0 {
		return nil
	}
	wp := workpool.New(c.cfg.Concurrency)
	for _, absFilename := range absFiles {
		var absName = absFilename
		wp.Do(func() error {
			var err error
			for i := 0; i < c.cfg.TryMax; i++ {
				start := time.Now()
				err = c.uploadToLinNode(absName)
				diff := time.Now().Sub(start)
				if diff < time.Second {
					time.Sleep(time.Second - diff)
				}
				if err == nil {
					break
				}
			}
			return err
		})
	}
	return wp.Wait()
}

func (c *Client) UploadRaw(key, contentType string, body io.Reader) (string, error) {
	if contentType == "" {
		contentType = defaultContentType
	}
	input := &s3manager.UploadInput{
		Bucket:      aws.String(c.cfg.Bucket),
		Key:         aws.String(key),
		ACL:         aws.String("public-read"),
		Body:        body,
		ContentType: aws.String(contentType),
	}
	result, err := s3manager.NewUploader(c.session, func(uploader *s3manager.Uploader) {
		uploader.PartSize = 1 * G
	}).Upload(input)
	if err != nil {
		log.Info("UpdateRaw Upload Err", log.Any("key", key), log.E(err))
		return "", err
	}
	log.Info("UpdateRaw Upload OK", log.Any("key", key), log.Any("Location", result.Location))
	return "", nil
}
func (c *Client) Upload(req UploadReq) error {
	if !req.Valid() {
		return errors.New("invalid upload request")
	}
	if req.ContentType == "" {
		req.ContentType = defaultContentType
		ext := path.Ext(req.Key)
		if ext != "" {
			req.ContentType = mime.TypeByExtension(ext)
		}
	}
	var err error
	for i := 0; i < defaultTryMax; i++ {
		buf := bytes.NewReader(req.Data)
		input := &s3.PutObjectInput{
			Bucket:      aws.String(c.cfg.Bucket),
			Key:         aws.String(req.Key),
			ACL:         aws.String("public-read"),
			Body:        buf,
			ContentType: aws.String(req.ContentType),
		}

		_, err := s3.New(c.session).PutObject(input)
		if err == nil {
			log.Info("Linode Upload OK", log.Any("key", req.Key))
			return nil
		}
		time.Sleep(time.Millisecond) //出让CPU
		log.Error("Linode Upload Try err", log.Any("key", req.Key), log.Any("tryCount", i+1), log.E(err))
	}
	log.Error("Linode Upload err", log.Any("key", req.Key), log.E(err))
	return err
}

// 多线程调用
func (m *Client) uploadToLinNodeWithKey(fileName, key string) error {
	f, err := os.OpenFile(fileName, os.O_RDONLY, 0644)
	if err != nil {
		log.Info("uploadToLinNodeWithKey OpenFile", log.Any("fileName", fileName), log.E(err))
		return err
	}
	defer func() { _ = f.Close() }()
	ext := filepath.Ext(fileName)
	contentType := ""
	if ext != "" {
		contentType = mime.TypeByExtension(ext)
	}
	_, err = m.UploadRaw(key, contentType, f)
	return err
}

func (m *Client) uploadToLinNode(absFilename string) error {
	key := strings.TrimPrefix(absFilename, fconst.FsDir)
	return m.uploadToLinNodeWithKey(absFilename, key)
}