Merge pull request #7 from xxxsen/xxxsen/feature/store_impl

Xxxsen/feature/store impl

Author: xxxsen

capture/capture.go

@@ -203,7 +203,7 @@ func (c *Capture) doNaming(ctx context.Context, fc *model.FileContext) error {
 
 func (c *Capture) doSaveData(ctx context.Context, fc *model.FileContext) error {
 	//保存元数据并将影片移入指定目录
-	if err := c.saveMediaData(fc); err != nil {
+	if err := c.saveMediaData(ctx, fc); err != nil {
 		return fmt.Errorf("save meta data failed, err:%w", err)
 	}
 	return nil
@@ -284,7 +284,7 @@ func (c *Capture) renameMetaField(fc *model.FileContext) error {
 	return nil
 }
 
-func (c *Capture) saveMediaData(fc *model.FileContext) error {
+func (c *Capture) saveMediaData(ctx context.Context, fc *model.FileContext) error {
 	images := make([]*model.File, 0, len(fc.Meta.SampleImages)+2)
 	if fc.Meta.Cover != nil {
 		images = append(images, fc.Meta.Cover)
@@ -297,7 +297,7 @@ func (c *Capture) saveMediaData(fc *model.FileContext) error {
 		target := filepath.Join(fc.SaveDir, image.Name)
 		logger := logutil.GetLogger(context.Background()).With(zap.String("image", image.Name), zap.String("key", image.Key), zap.String("target", target))
 
-		data, err := store.GetDefault().GetData(image.Key)
+		data, err := store.GetData(ctx, image.Key)
 		if err != nil {
 			logger.Error("read image data failed", zap.Error(err))
 			return err

hasher/hasher.go

@@ -0,0 +1,12 @@
+package hasher
+
+import (
+	"crypto/md5"
+	"encoding/hex"
+)
+
+func ToMD5(in string) string {
+	h := md5.New()
+	_, _ = h.Write([]byte(in))
+	return hex.EncodeToString(h.Sum(nil))
+}

main.go

@@ -38,9 +38,7 @@ func main() {
 	if err := precheckDir(c); err != nil {
 		logkit.Fatal("precheck dir failed", zap.Error(err))
 	}
-	if err := store.Init(filepath.Join(c.DataDir, "cache")); err != nil {
-		logkit.Fatal("init store failed", zap.Error(err))
-	}
+	store.SetStorage(store.NewDiskStorage(filepath.Join(c.DataDir, "cache")))
 	if err := translator.Init(); err != nil {
 		logkit.Error("init translater failed", zap.Error(err))
 	}

processor/handler/image_transcode_handler.go

@@ -40,28 +40,22 @@ func (p *imageTranscodeHandler) transcode(ctx context.Context, name string, f *m
 		return nil
 	}
 	logger = logger.With(zap.String("key", f.Key))
-	data, err := store.GetDefault().GetData(f.Key)
-	if err != nil {
-		logger.Debug("read key data failed", zap.Error(err))
-		return f //不丢弃, 后续处理的时候报错, 方便发现问题
-	}
-	raw, err := image.TranscodeToJpeg(data)
-	if err != nil && strings.Contains(err.Error(), "luma/chroma subsampling ratio") && ffmpeg.IsFFMpegEnabled() {
-		data, err = ffmpeg.ConvertToYuv420pJpegFromBytes(ctx, data)
-		if err != nil {
-			logger.Error("use ffmpeg to correct invalid image data failed", zap.Error(err))
-			return nil
+
+	key, err := store.AnonymousDataRewrite(ctx, f.Key, func(ctx context.Context, data []byte) ([]byte, error) {
+		raw, err := image.TranscodeToJpeg(data)
+		if err != nil && strings.Contains(err.Error(), "luma/chroma subsampling ratio") && ffmpeg.IsFFMpegEnabled() {
+			data, err = ffmpeg.ConvertToYuv420pJpegFromBytes(ctx, data)
+			if err != nil {
+				logger.Error("use ffmpeg to correct invalid image data failed", zap.Error(err))
+				return nil, err
+			}
+			raw, err = image.TranscodeToJpeg(data)
 		}
-		raw, err = image.TranscodeToJpeg(data)
-	}
-	if err != nil {
-		logger.Error("unable to convert image to jpeg format", zap.Error(err))
-		return nil
-	}
-	key, err := store.GetDefault().Put(raw)
+		return raw, err
+	})
 	if err != nil {
-		logger.Error("store transcoded image data failed", zap.Error(err))
-		return f //
+		logger.Error("transcoded image data failed", zap.Error(err))
+		return nil //
 	}
 	logger.Debug("transcode image succ", zap.String("new_key", key))
 	f.Key = key

processor/handler/poster_crop_handler.go

@@ -46,15 +46,9 @@ func (c *posterCropHandler) Handle(ctx context.Context, fc *model.FileContext) e
 	if fc.Number.GetIsUncensorMovie() && face.IsFaceRecognizeEnabled() { //如果为步兵, 则使用人脸识别(当然, 只有该特性能用的情况下才启用)
 		cutter = c.wrapCutImageWithFaceRec(ctx, image.CutCensoredImageFromBytes)
 	}
-	data, err := store.GetDefault().GetData(fc.Meta.Cover.Key)
-	if err != nil {
-		return fmt.Errorf("get cover data failed, err:%w, key:%s", err, fc.Meta.Cover.Key)
-	}
-	res, err := cutter(data)
-	if err != nil {
-		return fmt.Errorf("cut poster image failed, err:%w", err)
-	}
-	key, err := store.GetDefault().Put(res)
+	key, err := store.AnonymousDataRewrite(ctx, fc.Meta.Cover.Key, func(ctx context.Context, data []byte) ([]byte, error) {
+		return cutter(data)
+	})
 	if err != nil {
 		return fmt.Errorf("save cutted poster data failed, err:%w", err)
 	}

processor/handler/watermark_handler.go

@@ -17,10 +17,6 @@ func (h *watermark) Handle(ctx context.Context, fc *model.FileContext) error {
 	if fc.Meta.Poster == nil || len(fc.Meta.Poster.Key) == 0 {
 		return nil
 	}
-	data, err := store.GetDefault().GetData(fc.Meta.Poster.Key)
-	if err != nil {
-		return fmt.Errorf("load poster key failed, key:%s", fc.Meta.Poster.Key)
-	}
 	tags := make([]image.Watermark, 0, 5)
 	if fc.Number.GetIs4K() {
 		tags = append(tags, image.WM4K)
@@ -38,11 +34,9 @@ func (h *watermark) Handle(ctx context.Context, fc *model.FileContext) error {
 		logutil.GetLogger(ctx).Debug("no watermark tag found, skip watermark proc")
 		return nil
 	}
-	newData, err := image.AddWatermarkFromBytes(data, tags)
-	if err != nil {
-		return fmt.Errorf("add watermark failed, err:%w", err)
-	}
-	key, err := store.GetDefault().Put(newData)
+	key, err := store.AnonymousDataRewrite(ctx, fc.Meta.Poster.Key, func(ctx context.Context, data []byte) ([]byte, error) {
+		return image.AddWatermarkFromBytes(data, tags)
+	})
 	if err != nil {
 		return fmt.Errorf("save watermarked image failed, err:%w", err)
 	}

searcher/default_searcher.go

@@ -2,13 +2,12 @@ package searcher
 
 import (
 	"context"
-	"crypto/md5"
-	"encoding/hex"
 	"fmt"
 	"net/http"
 	"strings"
 	"time"
 	"yamdc/client"
+	"yamdc/hasher"
 	"yamdc/model"
 	"yamdc/number"
 	"yamdc/searcher/plugin"
@@ -230,8 +229,8 @@ func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []st
 			continue
 		}
 		logger := logutil.GetLogger(context.Background()).With(zap.String("url", url))
-		key := p.buildURLCacheKey(url)
-		if store.GetDefault().IsCacheExist(key) {
+		key := hasher.ToMD5(url)
+		if ok, _ := store.IsDataExist(ctx.GetContext(), key); ok {
 			rs[url] = key
 			continue
 		}
@@ -240,7 +239,7 @@ func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []st
 			logger.Error("fetch image data failed", zap.Error(err))
 			continue
 		}
-		err = store.GetDefault().PutWithNamingKey(key, data)
+		err = store.PutData(ctx.GetContext(), key, data)
 		if err != nil {
 			logger.Error("put image data to store failed", zap.Error(err))
 		}
@@ -249,14 +248,6 @@ func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []st
 	return rs
 }
 
-func (p *DefaultSearcher) buildURLCacheKey(url string) string {
-	h := md5.New()
-	_, _ = h.Write([]byte(url))
-	hashsum := hex.EncodeToString(h.Sum(nil))
-	key := fmt.Sprintf("avc:search:cache:url:%s:%s", p.name, hashsum)
-	return key
-}
-
 func (p *DefaultSearcher) fetchImageData(ctx *plugin.PluginContext, url string) ([]byte, error) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {

store/disk_storage.go

@@ -0,0 +1,53 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"yamdc/hasher"
+)
+
+type diskStorage struct {
+	dir string
+}
+
+func NewDiskStorage(dir string) IStorage {
+	return &diskStorage{dir: dir}
+}
+
+func (s *diskStorage) generateStorePath(key string) string {
+	save := hasher.ToMD5(key)
+	p1 := save[:2]
+	p2 := save[2:4]
+	p3 := save[4:6]
+	return filepath.Join(s.dir, p1, p2, p3, save)
+}
+
+func (s *diskStorage) GetData(ctx context.Context, key string) ([]byte, error) {
+	return os.ReadFile(s.generateStorePath(key))
+}
+
+func (s *diskStorage) PutData(ctx context.Context, key string, value []byte) error {
+	p := s.generateStorePath(key)
+	dir := filepath.Dir(p)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("create dir failed, err:%w", err)
+	}
+	if err := os.WriteFile(p, value, 0644); err != nil {
+		return fmt.Errorf("write data failed, err:%w", err)
+	}
+	return nil
+}
+
+func (s *diskStorage) IsDataExist(ctx context.Context, key string) (bool, error) {
+	p := s.generateStorePath(key)
+	_, err := os.Stat(p)
+	if err == nil {
+		return true, nil
+	}
+	if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}

store/disk_storage_test.go

@@ -0,0 +1,31 @@
+package store
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDiskStorage(t *testing.T) {
+	dir := filepath.Join(os.TempDir(), "storage")
+	st := NewDiskStorage(dir)
+	ctx := context.Background()
+	key := "aaa"
+	value := []byte("hello world")
+	err := st.PutData(ctx, key, value)
+	assert.NoError(t, err)
+	exist, err := st.IsDataExist(ctx, key)
+	assert.NoError(t, err)
+	assert.True(t, exist)
+	data, err := st.GetData(ctx, key)
+	assert.NoError(t, err)
+	assert.Equal(t, value, data)
+
+	//check not exist
+	exist, err = st.IsDataExist(ctx, "hello")
+	assert.NoError(t, err)
+	assert.False(t, exist)
+}

store/storage.go

@@ -0,0 +1,61 @@
+package store
+
+import (
+	"context"
+
+	"github.com/google/uuid"
+)
+
+type DataRewriteFunc func(ctx context.Context, data []byte) ([]byte, error)
+
+type IStorage interface {
+	GetData(ctx context.Context, key string) ([]byte, error)
+	PutData(ctx context.Context, key string, value []byte) error
+	IsDataExist(ctx context.Context, key string) (bool, error)
+}
+
+var defaultInst IStorage
+
+func SetStorage(impl IStorage) {
+	defaultInst = impl
+}
+
+func getDefaultInst() IStorage {
+	return defaultInst
+}
+
+func PutData(ctx context.Context, key string, value []byte) error {
+	return getDefaultInst().PutData(ctx, key, value)
+}
+
+func AnonymousPutData(ctx context.Context, value []byte) (string, error) {
+	key := uuid.NewString()
+	if err := PutData(ctx, key, value); err != nil {
+		return "", err
+	}
+	return key, nil
+}
+
+func GetData(ctx context.Context, key string) ([]byte, error) {
+	return getDefaultInst().GetData(ctx, key)
+}
+
+func IsDataExist(ctx context.Context, key string) (bool, error) {
+	return getDefaultInst().IsDataExist(ctx, key)
+}
+
+func AnonymousDataRewrite(ctx context.Context, key string, fn DataRewriteFunc) (string, error) {
+	raw, err := GetData(ctx, key)
+	if err != nil {
+		return key, err
+	}
+	newData, err := fn(ctx, raw)
+	if err != nil {
+		return key, err
+	}
+	newKey, err := AnonymousPutData(ctx, newData)
+	if err != nil {
+		return key, err
+	}
+	return newKey, nil
+}