Merge pull request #11 from xxxsen/xxxsen/feature/rewrite_plugin_api

Xxxsen/feature/rewrite plugin api

capture/capture.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+	"yamdc/envflag"
 	"yamdc/model"
 	"yamdc/nfo"
 	"yamdc/number"
@@ -320,7 +321,7 @@ func (c *Capture) saveMediaData(ctx context.Context, fc *model.FileContext) erro
 }
 
 func (c *Capture) moveMovie(fc *model.FileContext, src string, dst string) error {
-	if c.c.LinkMode {
+	if envflag.IsEnableLinkMode() {
 		return c.moveMovieByLink(fc, src, dst)
 	}
 	return c.moveMovieDirect(fc, src, dst)

capture/config.go

@@ -24,7 +24,6 @@ type config struct {
 	SaveDir           string
 	Naming            string
 	ExtraMediaExtList []string
-	LinkMode          bool
 }
 
 type Option func(c *config)
@@ -59,12 +58,6 @@ func WithNamingRule(r string) Option {
 	}
 }
 
-func WithEnableLinkMode(v bool) Option {
-	return func(c *config) {
-		c.LinkMode = v
-	}
-}
-
 func WithExtraMediaExtList(lst []string) Option {
 	return func(c *config) {
 		c.ExtraMediaExtList = lst

envflag/envflag.go

@@ -0,0 +1,33 @@
+package envflag
+
+import (
+	"github.com/kelseyhightower/envconfig"
+)
+
+var defaultInst = &EnvFlag{}
+
+type EnvFlag struct {
+	EnableSearchMetaCache bool `envconfig:"enable_search_meta_cache" default:"true"`
+	EnableLinkMode        bool `envconfig:"enable_link_mode"`
+}
+
+func GetFlag() *EnvFlag {
+	return defaultInst
+}
+
+func Init() error {
+	fg := &EnvFlag{}
+	if err := envconfig.Process("exec_flag", fg); err != nil {
+		return err
+	}
+	defaultInst = fg
+	return nil
+}
+
+func IsEnableSearchMetaCache() bool {
+	return GetFlag().EnableSearchMetaCache
+}
+
+func IsEnableLinkMode() bool {
+	return GetFlag().EnableLinkMode
+}

go.mod

@@ -31,6 +31,7 @@ require (
 	github.com/google/pprof v0.0.0-20240910150728-a0b0bb1d4134 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
+	github.com/kelseyhightower/envconfig v1.4.0
 	github.com/klauspost/compress v1.17.9 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/onsi/ginkgo/v2 v2.20.2 // indirect

go.sum

@@ -43,6 +43,8 @@ github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+l
 github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
 github.com/imroc/req/v3 v3.48.0 h1:IYuMGetuwLzOOTzDCquDqs912WNwpsPK0TBXWPIvoqg=
 github.com/imroc/req/v3 v3.48.0/go.mod h1:weam9gmyb00QnOtu6HXSnk44dNFkIUQb5QdMx13FeUU=
+github.com/kelseyhightower/envconfig v1.4.0 h1:Im6hONhd3pLkfDFsbRgu68RDNkGF1r3dvMUtDTo2cv8=
+github.com/kelseyhightower/envconfig v1.4.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg=
 github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=

main.go

@@ -11,6 +11,7 @@ import (
 	"yamdc/client"
 	"yamdc/config"
 	"yamdc/dependency"
+	"yamdc/envflag"
 	"yamdc/face"
 	"yamdc/face/goface"
 	"yamdc/face/pigo"
@@ -26,8 +27,8 @@ import (
 	"github.com/xxxsen/common/logutil"
 	"go.uber.org/zap"
 
-	"yamdc/searcher/plugin"
-	_ "yamdc/searcher/plugin/airav"
+	"yamdc/searcher/plugin/factory"
+	_ "yamdc/searcher/plugin/register"
 )
 
 var conf = flag.String("config", "./config.json", "config file")
@@ -47,14 +48,20 @@ func main() {
 		logkit.Fatal("ensure dependencies failed", zap.Error(err))
 	}
 	logkit.Info("check dependencies finish...")
+
+	if err := envflag.Init(); err != nil {
+		logkit.Fatal("init envflag failed", zap.Error(err))
+	}
+	logkit.Info("read env flags", zap.Any("flag", *envflag.GetFlag()))
+
 	store.SetStorage(store.MustNewSqliteStorage(filepath.Join(c.DataDir, "cache", "cache.db")))
 	if err := translator.Init(); err != nil {
 		logkit.Error("init translater failed", zap.Error(err))
 	}
 	if err := initFace(filepath.Join(c.DataDir, "models")); err != nil {
 		logkit.Error("init face recognizer failed", zap.Error(err))
 	}
-	logkit.Info("support plugins", zap.Strings("plugins", plugin.Plugins()))
+	logkit.Info("support plugins", zap.Strings("plugins", factory.Plugins()))
 	logkit.Info("support handlers", zap.Strings("handlers", handler.Handlers()))
 	logkit.Info("current use plugins", zap.Strings("plugins", c.Plugins))
 	for _, ct := range c.CategoryPlugins {
@@ -104,7 +111,6 @@ func buildCapture(c *config.Config, ss []searcher.ISearcher, catSs map[number.Ca
 		capture.WithSaveDir(c.SaveDir),
 		capture.WithSeacher(searcher.NewCategorySearcher(ss, catSs)),
 		capture.WithProcessor(processor.NewGroup(ps)),
-		capture.WithEnableLinkMode(c.SwitchConfig.EnableLinkMode),
 		capture.WithExtraMediaExtList(c.ExtraMediaExts),
 	)
 	return capture.New(opts...)
@@ -129,7 +135,7 @@ func buildSearcher(plgs []string, m map[string]interface{}) ([]searcher.ISearche
 		if !ok {
 			args = struct{}{}
 		}
-		plg, err := plugin.CreatePlugin(name, args)
+		plg, err := factory.CreatePlugin(name, args)
 		if err != nil {
 			return nil, fmt.Errorf("create plugin failed, name:%s, err:%w", name, err)
 		}

processor/handler/tag_padder_handler.go

@@ -29,12 +29,26 @@ func (h *tagPadderHandler) generateNumberPrefixTag(fc *model.FileContext) (strin
 	return sb.String(), true
 }
 
+func (h *tagPadderHandler) rewriteOrAppendTag(fc *model.AvMeta, tagname string) {
+	isContained := false
+	for idx, item := range fc.Genres {
+		if strings.EqualFold(item, tagname) {
+			fc.Genres[idx] = tagname
+			isContained = true
+		}
+	}
+	if isContained {
+		return
+	}
+	fc.Genres = append(fc.Genres, tagname)
+}
+
 func (h *tagPadderHandler) Handle(ctx context.Context, fc *model.FileContext) error {
 	//提取番号特有的tag
 	fc.Meta.Genres = append(fc.Meta.Genres, fc.Number.GenerateTags()...)
 	//提取番号前缀作为tag
 	if tag, ok := h.generateNumberPrefixTag(fc); ok {
-		fc.Meta.Genres = append(fc.Meta.Genres, tag)
+		h.rewriteOrAppendTag(fc.Meta, tag)
 	}
 	fc.Meta.Genres = utils.DedupStringList(fc.Meta.Genres)
 	return nil

searcher/default_searcher.go

@@ -7,10 +7,12 @@ import (
 	"strings"
 	"time"
 	"yamdc/client"
+	"yamdc/envflag"
 	"yamdc/hasher"
 	"yamdc/model"
 	"yamdc/number"
-	"yamdc/searcher/plugin"
+	"yamdc/searcher/plugin/api"
+	"yamdc/searcher/plugin/meta"
 	"yamdc/store"
 	"yamdc/useragent"
 
@@ -25,26 +27,26 @@ const (
 type DefaultSearcher struct {
 	name    string
 	ua      string
-	invoker plugin.HTTPInvoker
-	plg     plugin.IPlugin
+	invoker api.HTTPInvoker
+	plg     api.IPlugin
 }
 
-func MustNewDefaultSearcher(name string, plg plugin.IPlugin) ISearcher {
+func MustNewDefaultSearcher(name string, plg api.IPlugin) ISearcher {
 	s, err := NewDefaultSearcher(name, plg)
 	if err != nil {
 		panic(err)
 	}
 	return s
 }
 
-func defaultInvoker() plugin.HTTPInvoker {
+func defaultInvoker() api.HTTPInvoker {
 	basicClient := client.NewClient()
-	return func(ctx *plugin.PluginContext, req *http.Request) (*http.Response, error) {
+	return func(ctx context.Context, req *http.Request) (*http.Response, error) {
 		return basicClient.Do(req)
 	}
 }
 
-func NewDefaultSearcher(name string, plg plugin.IPlugin) (ISearcher, error) {
+func NewDefaultSearcher(name string, plg api.IPlugin) (ISearcher, error) {
 	invoker := plg.OnHTTPClientInit()
 	if invoker == nil {
 		invoker = defaultInvoker()
@@ -72,7 +74,7 @@ func (p *DefaultSearcher) setDefaultHttpOptions(req *http.Request) error {
 	return nil
 }
 
-func (p *DefaultSearcher) decorateRequest(ctx *plugin.PluginContext, req *http.Request) error {
+func (p *DefaultSearcher) decorateRequest(ctx context.Context, req *http.Request) error {
 	if err := p.plg.OnDecorateRequest(ctx, req); err != nil {
 		return err
 	}
@@ -82,7 +84,7 @@ func (p *DefaultSearcher) decorateRequest(ctx *plugin.PluginContext, req *http.R
 	return nil
 }
 
-func (p *DefaultSearcher) decorateImageRequest(ctx *plugin.PluginContext, req *http.Request) error {
+func (p *DefaultSearcher) decorateImageRequest(ctx context.Context, req *http.Request) error {
 	if err := p.plg.OnDecorateMediaRequest(ctx, req); err != nil {
 		return err
 	}
@@ -92,21 +94,21 @@ func (p *DefaultSearcher) decorateImageRequest(ctx *plugin.PluginContext, req *h
 	return nil
 }
 
-func (p *DefaultSearcher) invokeHTTPRequest(ctx *plugin.PluginContext, req *http.Request) (*http.Response, error) {
+func (p *DefaultSearcher) invokeHTTPRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
 	if err := p.decorateRequest(ctx, req); err != nil {
 		return nil, fmt.Errorf("decorate request failed, err:%w", err)
 	}
 	return p.invoker(ctx, req)
 }
 
-func (p *DefaultSearcher) onRetriveData(ctx context.Context, pctx *plugin.PluginContext, req *http.Request, number *number.Number) ([]byte, error) {
+func (p *DefaultSearcher) onRetriveData(ctx context.Context, req *http.Request, number *number.Number) ([]byte, error) {
 	key := p.name + ":" + number.GetNumberID()
-	return store.LoadData(ctx, key, defaultPageSearchCacheExpire, func() ([]byte, error) {
-		rsp, err := p.plg.OnHandleHTTPRequest(pctx, p.invokeHTTPRequest, req)
+	dataLoader := func() ([]byte, error) {
+		rsp, err := p.plg.OnHandleHTTPRequest(ctx, p.invokeHTTPRequest, req)
 		if err != nil {
 			return nil, fmt.Errorf("do request failed, err:%w", err)
 		}
-		isSearchSucc, err := p.plg.OnPrecheckResponse(pctx, req, rsp)
+		isSearchSucc, err := p.plg.OnPrecheckResponse(ctx, req, rsp)
 		if err != nil {
 			return nil, fmt.Errorf("precheck responnse failed, err:%w", err)
 		}
@@ -122,28 +124,31 @@ func (p *DefaultSearcher) onRetriveData(ctx context.Context, pctx *plugin.Plugin
 			return nil, fmt.Errorf("read body failed, err:%w", err)
 		}
 		return data, nil
-	})
+	}
+	if !envflag.IsEnableSearchMetaCache() {
+		return dataLoader()
+	}
+	return store.LoadData(ctx, key, defaultPageSearchCacheExpire, dataLoader)
 }
 
 func (p *DefaultSearcher) Search(ctx context.Context, number *number.Number) (*model.AvMeta, bool, error) {
-	pctx := plugin.NewPluginContext(ctx)
-	pctx.SetNumberInfo(number)
-	ok, err := p.plg.OnPrecheckRequest(pctx, number)
+	ctx = meta.SetNumberId(ctx, number.GetNumberID())
+	ok, err := p.plg.OnPrecheckRequest(ctx, number)
 	if err != nil {
 		return nil, false, fmt.Errorf("precheck failed, err:%w", err)
 	}
 	if !ok {
 		return nil, false, nil
 	}
-	req, err := p.plg.OnMakeHTTPRequest(pctx, number)
+	req, err := p.plg.OnMakeHTTPRequest(ctx, number)
 	if err != nil {
 		return nil, false, fmt.Errorf("make http request failed, err:%w", err)
 	}
-	data, err := p.onRetriveData(ctx, pctx, req, number)
+	data, err := p.onRetriveData(ctx, req, number)
 	if err != nil {
 		return nil, false, err
 	}
-	meta, decodeSucc, err := p.plg.OnDecodeHTTPData(pctx, data)
+	meta, decodeSucc, err := p.plg.OnDecodeHTTPData(ctx, data)
 	if err != nil {
 		return nil, false, fmt.Errorf("decode http data failed, err:%w", err)
 	}
@@ -153,7 +158,7 @@ func (p *DefaultSearcher) Search(ctx context.Context, number *number.Number) (*m
 	//重建不规范的元数据
 	p.fixMeta(req, meta)
 	//将远程数据保存到本地, 并替换文件key
-	p.storeImageData(pctx, meta)
+	p.storeImageData(ctx, meta)
 	if err := p.verifyMeta(meta); err != nil {
 		logutil.GetLogger(ctx).Error("verify meta not pass, treat as not found", zap.Error(err), zap.String("plugin", p.name))
 		return nil, false, nil
@@ -204,7 +209,7 @@ func (p *DefaultSearcher) fixSingleURL(req *http.Request, input *string, prefix
 	}
 }
 
-func (p *DefaultSearcher) storeImageData(ctx *plugin.PluginContext, in *model.AvMeta) {
+func (p *DefaultSearcher) storeImageData(ctx context.Context, in *model.AvMeta) {
 	images := make([]string, 0, len(in.SampleImages)+2)
 	if in.Cover != nil {
 		images = append(images, in.Cover.Name)
@@ -237,15 +242,15 @@ func (p *DefaultSearcher) storeImageData(ctx *plugin.PluginContext, in *model.Av
 	in.SampleImages = rebuildSampleList
 }
 
-func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []string) map[string]string {
+func (p *DefaultSearcher) saveRemoteURLData(ctx context.Context, urls []string) map[string]string {
 	rs := make(map[string]string, len(urls))
 	for _, url := range urls {
 		if len(url) == 0 {
 			continue
 		}
 		logger := logutil.GetLogger(context.Background()).With(zap.String("url", url))
 		key := hasher.ToSha1(url)
-		if ok, _ := store.IsDataExist(ctx.GetContext(), key); ok {
+		if ok, _ := store.IsDataExist(ctx, key); ok {
 			rs[url] = key
 			continue
 		}
@@ -254,7 +259,7 @@ func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []st
 			logger.Error("fetch image data failed", zap.Error(err))
 			continue
 		}
-		err = store.PutData(ctx.GetContext(), key, data)
+		err = store.PutData(ctx, key, data)
 		if err != nil {
 			logger.Error("put image data to store failed", zap.Error(err))
 		}
@@ -263,7 +268,7 @@ func (p *DefaultSearcher) saveRemoteURLData(ctx *plugin.PluginContext, urls []st
 	return rs
 }
 
-func (p *DefaultSearcher) fetchImageData(ctx *plugin.PluginContext, url string) ([]byte, error) {
+func (p *DefaultSearcher) fetchImageData(ctx context.Context, url string) ([]byte, error) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
 		return nil, fmt.Errorf("make request for url:%s failed, err:%w", url, err)

searcher/plugin/api/api.go

@@ -0,0 +1,21 @@
+package api
+
+import (
+	"context"
+	"net/http"
+	"yamdc/model"
+	"yamdc/number"
+)
+
+type HTTPInvoker func(ctx context.Context, req *http.Request) (*http.Response, error)
+
+type IPlugin interface {
+	OnHTTPClientInit() HTTPInvoker
+	OnPrecheckRequest(ctx context.Context, number *number.Number) (bool, error)
+	OnMakeHTTPRequest(ctx context.Context, number *number.Number) (*http.Request, error)
+	OnDecorateRequest(ctx context.Context, req *http.Request) error
+	OnHandleHTTPRequest(ctx context.Context, invoker HTTPInvoker, req *http.Request) (*http.Response, error)
+	OnPrecheckResponse(ctx context.Context, req *http.Request, rsp *http.Response) (bool, error)
+	OnDecodeHTTPData(ctx context.Context, data []byte) (*model.AvMeta, bool, error)
+	OnDecorateMediaRequest(ctx context.Context, req *http.Request) error
+}