Files
kenaito-media-picker/internal/organize/organize.go
T
odboy f915a9b1ac feat: 扫描整理并行化、WebSocket 实时进度与统计看板
- 移动阶段并行处理文件,按原子计数上报累计进度
- 扫描遍历阶段实时上报已发现文件数,界面不再卡在 0/0
- 新增 /api/scan/ws WebSocket 实时推送扫描状态;前端断线指数退避自动重连,并以低频轮询兜底
- 新增统计看板与媒体库筛选/排序,目录选择器路径不存在时回退到最近可用父目录
- 存放目录支持自定义并持久化,默认输出到 KMP已整理
2025-05-08 21:16:20 +08:00

441 lines
13 KiB
Go

// Package organize 编排整理流程:探测 → 归类 → 收集 bucket → 一次性移动并记录去向。
package organize
import (
"fmt"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"sync/atomic"
"kenaito-media-picker/internal/classify"
"kenaito-media-picker/internal/digest"
"kenaito-media-picker/internal/mover"
"kenaito-media-picker/internal/probe"
"kenaito-media-picker/internal/scanner"
"kenaito-media-picker/internal/store"
)
// 特殊归类目录名
const (
dirImagePSD = "图片_PSD"
dirImageErr = "图片_读取异常"
dirImageIgnored = "图片_已忽略"
dirVideoErr = "视频_读取异常"
dirVideoIgnored = "视频_已忽略"
dirAudio = "音频"
dirDuplicate = "重复文件"
)
// Organizer 整理器:按目标目录收集文件,处理完成后统一建目录并移动
type Organizer struct {
pickerDir string
ffprobePath string
store *store.Store
mover *mover.Mover
buckets map[string][]string // 目标目录 → 文件路径
order []string // 目标目录的创建顺序(保证输出稳定)
seenMd5 map[string]bool // 本次扫描已入库的 MD5(同一扫描内去重)
onProgress func(done, total int)
pending []*store.MediaInfo // 待批量入库的媒体记录
pendingDup []*store.MediaInfo // 待批量入库的重复文件记录
}
// New 创建整理器;store 可为 nil(数据库初始化失败时仅不记录信息)
func New(pickerDir, ffprobePath string, st *store.Store) *Organizer {
return &Organizer{
pickerDir: pickerDir,
ffprobePath: ffprobePath,
store: st,
mover: &mover.Mover{},
buckets: make(map[string][]string),
seenMd5: make(map[string]bool),
}
}
// SetProgress 设置进度回调(done/total 为已处理文件数与总数)
func (o *Organizer) SetProgress(fn func(done, total int)) {
o.onProgress = fn
}
// report 上报进度
func (o *Organizer) report(done, total int) {
if o.onProgress != nil && total > 0 {
o.onProgress(done, total)
}
}
// workerCount 并行探测的 worker 数量:不超过 CPU 核数,且限制在 1~8 之间
func workerCount() int {
n := runtime.NumCPU()
if n < 1 {
return 1
}
if n > 8 {
return 8
}
return n
}
// runParallel 用固定 worker 数并行处理切片,按原顺序收集结果。
// onProgress 在每次完成一个任务后调用(done 为已完成任务数,并发安全)。
// 返回 results[i] 对应 items[i] 的处理结果。
func runParallel[T any, R any](items []T, fn func(T) R, onProgress func(done int)) []R {
results := make([]R, len(items))
if len(items) == 0 {
return results
}
workers := workerCount()
if workers > len(items) {
workers = len(items)
}
var wg sync.WaitGroup
var doneCount int32
jobs := make(chan int)
for w := 0; w < workers; w++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := range jobs {
results[i] = fn(items[i])
if onProgress != nil {
onProgress(int(atomic.AddInt32(&doneCount, 1)))
}
}
}()
}
for i := range items {
jobs <- i
}
close(jobs)
wg.Wait()
return results
}
// add 将文件归入某个子目录
func (o *Organizer) add(f scanner.File, subDir string) {
dir := filepath.Join(o.pickerDir, subDir)
if _, ok := o.buckets[dir]; !ok {
o.order = append(o.order, dir)
}
o.buckets[dir] = append(o.buckets[dir], f.Path)
}
// calcFileInfo 读取文件大小并计算 MD5(大文件优先用指纹复用已入库的 MD5,跳过全量读取)。
// store 为 nil 时直接返回空值(无需去重/入库)。
func (o *Organizer) calcFileInfo(f scanner.File) (md5Str, headMd5, tailMd5 string, size int64) {
if o.store == nil {
return "", "", "", 0
}
st, err := os.Stat(f.Path)
if err != nil {
fmt.Printf("=== 读取文件信息失败: %s, 错误: %v\n", f.Path, err)
return "", "", "", 0
}
size = st.Size()
if size > digest.FingerprintThreshold {
if head, tail, ferr := digest.CalcFileFingerprint(f.Path, size); ferr == nil {
headMd5, tailMd5 = head, tail
if existing, ok := o.store.FindMediaByFingerprint(size, head, tail); ok && existing.FileMd5 != "" {
md5Str = existing.FileMd5 // 指纹命中,复用MD5,跳过全量读取
}
}
}
if md5Str == "" {
if md5Str, err = digest.CalcFileMd5(f.Path); err != nil {
fmt.Printf("=== 计算MD5失败: %s, 错误: %v\n", f.Path, err)
return "", "", "", size
}
}
return md5Str, headMd5, tailMd5, size
}
// isDuplicate 判断文件是否重复:
// 1. 本次扫描已入库的 MD5 → 重复(同一扫描内去重)
// 2. 库中已有记录且其副本在磁盘上仍存在(DestPath)→ 重复
// 若库中副本已被移动/改名/删除,说明是同一文件被重新扫描,不算重复。
func (o *Organizer) isDuplicate(md5Str string) bool {
if md5Str == "" {
return false
}
if o.seenMd5[md5Str] {
return true
}
if o.store == nil {
return false
}
existing, ok := o.store.FindMediaByMd5(md5Str)
if !ok {
return false
}
if existing.DestPath != "" {
if _, err := os.Stat(existing.DestPath); err == nil {
return true
}
}
return false
}
// markSeen 记录本次扫描已入库的 MD5(去重用)
func (o *Organizer) markSeen(md5Str string) {
if md5Str != "" {
o.seenMd5[md5Str] = true
}
}
// record 收集媒体信息到待入库缓冲(不立即写库,ProcessImages/ProcessMedia 结束统一批量提交)
func (o *Organizer) record(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) {
if o.store == nil || md5Str == "" {
return
}
o.pending = append(o.pending, buildMediaInfo(f, mediaType, res, md5Str, headMd5, tailMd5, size))
}
// recordDuplicate 收集重复文件信息到待入库缓冲
func (o *Organizer) recordDuplicate(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) {
if o.store == nil || md5Str == "" {
return
}
o.pendingDup = append(o.pendingDup, buildMediaInfo(f, mediaType, res, md5Str, headMd5, tailMd5, size))
}
// buildMediaInfo 构造入库记录
func buildMediaInfo(f scanner.File, mediaType string, res *probe.ProbeResult, md5Str, headMd5, tailMd5 string, size int64) *store.MediaInfo {
codec := res.VideoCodec
if mediaType == "audio" {
codec = res.AudioCodec
}
format := res.Format
if format == "" && f.Ext != "" {
format = strings.TrimPrefix(f.Ext, ".")
}
return &store.MediaInfo{
FileMd5: md5Str,
FilePath: f.Path,
FileName: f.Name,
FileExt: f.Ext,
FileSize: size,
MimeType: f.Mime,
MediaType: mediaType,
Width: res.Width,
Height: res.Height,
Duration: res.Duration,
Codec: codec,
Format: format,
RawInfo: res.RawJSON,
HeadMd5: headMd5,
TailMd5: tailMd5,
}
}
// flushRecords 将缓冲的媒体/重复记录批量提交入库
func (o *Organizer) flushRecords() {
if o.store == nil {
return
}
if len(o.pending) > 0 {
if err := o.store.SaveMediaBatch(o.pending); err != nil {
fmt.Printf("=== 批量保存媒体信息失败: %v\n", err)
}
o.pending = nil
}
if len(o.pendingDup) > 0 {
if err := o.store.InsertDuplicateBatch(o.pendingDup); err != nil {
fmt.Printf("=== 批量保存重复文件信息失败: %v\n", err)
}
o.pendingDup = nil
}
}
// imageProbeResult 图片探测结果(并行阶段产出,串行阶段归类)
type imageProbeResult struct {
psd bool // PSD 特殊归类
unsupported bool // 非白名单且无 ffprobe 兜底,忽略
probeErr bool // 探测失败,归入异常
res *probe.ProbeResult
md5Str string
headMd5 string
tailMd5 string
size int64
darkScore float64
lightScore float64
}
// probeImage 并行探测一张图片:探测尺寸 + MD5 + 明暗分析(不含重复判断/归类/入库)。
// 注意:不访问 o.seenMd5 等共享可变状态,仅读 store(指纹复用),并行安全。
func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
r := imageProbeResult{}
if strings.EqualFold(f.Ext, ".psd") {
r.psd = true
return r
}
if !probe.IsSupportImage(f.Ext) && o.ffprobePath == "" {
r.unsupported = true
return r
}
res, err := probe.ProbeImage(f.Path, o.ffprobePath)
if err != nil {
r.probeErr = true
return r
}
r.res = res
r.md5Str, r.headMd5, r.tailMd5, r.size = o.calcFileInfo(f)
if !strings.EqualFold(f.Ext, ".gif") {
dark, light, berr := probe.AnalyzeBrightness(f.Path)
if berr != nil {
fmt.Printf("=== 明暗分析失败(按中性mid归类): %s \n", f.Path)
} else {
r.darkScore, r.lightScore = dark, light
}
}
return r
}
// ProcessImages 处理全部图片:并行探测(尺寸/MD5/明暗)→ 串行去重归类
func (o *Organizer) ProcessImages(images []scanner.File) {
success, fail, ignore, dup := 0, 0, 0, 0
total := len(images)
results := runParallel(images, o.probeImage, func(done int) {
o.report(done, total)
})
for i, r := range results {
f := images[i]
switch {
case r.psd:
o.add(f, dirImagePSD)
success++
case r.unsupported:
o.add(f, dirImageIgnored)
ignore++
case r.probeErr:
o.add(f, dirImageErr)
fail++
fmt.Printf("=== 异常图片: %s \n", f.Path)
case o.isDuplicate(r.md5Str):
o.recordDuplicate(f, "image", r.res, r.md5Str, r.headMd5, r.tailMd5, r.size)
o.add(f, dirDuplicate)
dup++
fmt.Printf("=== 重复图片: %s \n", f.Path)
default:
o.record(f, "image", r.res, r.md5Str, r.headMd5, r.tailMd5, r.size)
o.markSeen(r.md5Str)
o.add(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore))
success++
}
}
fmt.Printf("=== 图片处理完毕,成功: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", success, fail, ignore, dup)
}
// mediaProbeResult 媒体探测结果(并行阶段产出,串行阶段归类)
type mediaProbeResult struct {
res *probe.ProbeResult
err error
md5 string
head string
tail string
size int64
}
// probeMedia 并行探测一个媒体文件(ffprobe + MD5)。不访问共享可变状态。
func (o *Organizer) probeMedia(f scanner.File) mediaProbeResult {
res, err := probe.ProbeMediaFile(o.ffprobePath, f.Path)
if err != nil {
return mediaProbeResult{err: err}
}
if !res.HasVideo && !res.HasAudio {
return mediaProbeResult{res: res}
}
md5, head, tail, size := o.calcFileInfo(f)
return mediaProbeResult{res: res, md5: md5, head: head, tail: tail, size: size}
}
// ProcessMedia 处理媒体候选:并行 ffprobe 判定 → 串行去重归类
func (o *Organizer) ProcessMedia(media []scanner.File) {
video, audio, fail, ignore, dup := 0, 0, 0, 0, 0
total := len(media)
results := runParallel(media, o.probeMedia, func(done int) {
o.report(done, total)
})
for i, r := range results {
f := media[i]
if r.err != nil {
// 仅 MIME 明确为视频/音频的文件才移入异常/忽略目录;
// octet-stream/空 MIME 且 ffprobe 无法确认的原地忽略,不移动
if strings.HasPrefix(f.Mime, "video/") {
fail++
o.add(f, dirVideoErr)
fmt.Printf("=== 异常视频: %s \n", f.Path)
} else if scanner.IsExplicitMediaMime(f.Mime) {
ignore++
o.add(f, dirVideoIgnored)
} else {
ignore++
fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
}
continue
}
switch {
case r.res.HasVideo:
if o.isDuplicate(r.md5) {
o.recordDuplicate(f, "video", r.res, r.md5, r.head, r.tail, r.size)
dup++
o.add(f, dirDuplicate)
fmt.Printf("=== 重复媒体: %s \n", f.Path)
continue
}
video++
o.record(f, "video", r.res, r.md5, r.head, r.tail, r.size)
o.markSeen(r.md5)
o.add(f, classify.VideoDir(f.Ext, r.res.Width, r.res.Height))
case r.res.HasAudio:
if o.isDuplicate(r.md5) {
o.recordDuplicate(f, "audio", r.res, r.md5, r.head, r.tail, r.size)
dup++
o.add(f, dirDuplicate)
fmt.Printf("=== 重复媒体: %s \n", f.Path)
continue
}
audio++
o.record(f, "audio", r.res, r.md5, r.head, r.tail, r.size)
o.markSeen(r.md5)
o.add(f, dirAudio)
default:
// ffprobe 无音视频流:MIME 明确为媒体的归入忽略目录,其余原地不动
if scanner.IsExplicitMediaMime(f.Mime) {
ignore++
o.add(f, dirVideoIgnored)
} else {
ignore++
fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
}
}
}
fmt.Printf("=== 媒体处理完毕,视频: %d, 音频: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", video, audio, fail, ignore, dup)
}
// Flush 先将缓冲的媒体记录批量入库,再将收集到的文件移动到对应目录并回填归类去向。
// 移动期间按累计完成数上报进度。
func (o *Organizer) Flush() {
o.flushRecords()
total := 0
for _, dir := range o.order {
total += len(o.buckets[dir])
}
movedAll := 0
for _, dir := range o.order {
onProgress := func(done int) {
o.report(movedAll+done, total)
}
moved := o.mover.MoveToDir(o.buckets[dir], dir, onProgress)
movedAll += len(o.buckets[dir])
for src, dest := range moved {
if err := o.store.UpdateMediaDest(src, dir, dest); err != nil {
fmt.Printf("=== 更新媒体去向失败: %s, 错误: %v\n", src, err)
}
}
}
}