Skip to content

Commit

Permalink
Merge pull request #7 from xxxsen/xxxsen/feature/store_impl
Browse files Browse the repository at this point in the history
Xxxsen/feature/store impl
  • Loading branch information
xxxsen committed Aug 27, 2024
2 parents 1136abd + 74c7919 commit 611b1df
Show file tree
Hide file tree
Showing 12 changed files with 185 additions and 221 deletions.
6 changes: 3 additions & 3 deletions capture/capture.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
12 changes: 12 additions & 0 deletions hasher/hasher.go
Original file line number Diff line number Diff line change
@@ -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))
}
4 changes: 1 addition & 3 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
Expand Down
34 changes: 14 additions & 20 deletions processor/handler/image_transcode_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 3 additions & 9 deletions processor/handler/poster_crop_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
12 changes: 3 additions & 9 deletions processor/handler/watermark_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
}
Expand Down
17 changes: 4 additions & 13 deletions searcher/default_searcher.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
}
Expand All @@ -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))
}
Expand All @@ -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 {
Expand Down
53 changes: 53 additions & 0 deletions store/disk_storage.go
Original file line number Diff line number Diff line change
@@ -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
}
31 changes: 31 additions & 0 deletions store/disk_storage_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
61 changes: 61 additions & 0 deletions store/storage.go
Original file line number Diff line number Diff line change
@@ -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
}
Loading

0 comments on commit 611b1df

Please sign in to comment.