joybuywork/internal/plugins/export/task.go

package export

import (
	"context"
	"math"
	"os"
	"sync"
)

const (
	MaxRow = 1000000 // excel 最大行数
)

// RowData 单条数据类型
type RowData map[string]interface{}

// AfterFinishCustomize 完成后自定义处理函数
type AfterFinishCustomize func(fileName, fullPath string) error

type Task interface {
	GetTotal() int64                      //获取导出的总条数
	GetBaseTask() *BaseTask               //获取基础任务结构
	WriteData(data []RowData) (err error) //写入数据
	FileClose()                           //关闭文件
	GetFilename() string                  //获取文件名称
	GetFullPath() string                  //获取文件存储路径
}

// Done 执行导出任务，虽然现在没有错误
func Done(et Task) {
	//没有数据导出的时候直接任务结束
	t := et.GetBaseTask()
	if et.GetTotal() == 0 {
		t.taskEnd(nil)
		return
	}

	//监听步进完成
	go t.listenNotifySp()

	//监听数据加入
	go t.listenAppendData()

	//监听任务情况
	go func() {
		var taskErr error
		select {
		case <-t.finish: //正常完成触发
			break
		case <-t.ctx.Done(): //超时触发
			break
		case notifyErr, ok := <-t.err: //执行过程出现异常
			if ok {
				taskErr = notifyErr
			}
		}

		//任务结束后
		t.taskEnd(taskErr)
	}()

	return
}

type BaseTask struct {
	ctx    context.Context    //上下文
	cancel context.CancelFunc //关闭上下文
	err    chan error         //错误通知，必须要有错误才通知
	fr     int64              //当前写入的行数
	finish chan bool          //导出完成标识
	ap     chan []RowData     //需要导出的数据
	reply  chan Reply         //处理结果通道，文件导出完成后，会通知该通道返回处理的结果，调用AfterFinish()方法可以获取到通知
	ti     Task               //任务接口(ti=>taskInterface)
	fp     *FinishPart        //使用完成部分通讯
	pe     *Proceed           //从某一条数据之后继续导出
}

type Proceed struct {
	l         *sync.Mutex //写入进度锁
	filterNum int64       //先过滤多少条数据之后才开始计算
	nft       int64       //当前已经过滤了多少条数据
	begin     int64       //到当写入条数到这个点的时候才开始写入文件
}

type PartDone func(spn, nfp float64)

// FinishPart 完成部分导出的结构体
type FinishPart struct {
	l        *sync.Mutex  //写入进度锁
	spn      float64      //把整个导出文件分成多少份
	sp       int64        //每多少条后计算一次速率
	nfp      float64      //当前已导出的份数
	notifySp chan float64 //当前进度 / 分割数的值
	fn       PartDone     //进度完成要做什么操作
}

// Reply 处理结果结构体
type Reply struct {
	Error    error
	Filename string
	FilePath string
	Fr       int64
}

// Init 初始化基础任务
func (t *BaseTask) Init(ctx context.Context, ti Task, opts ...TaskOption) {
	//设置的数据
	for _, opt := range opts {
		opt(t)
	}

	//初始化任务值
	t.err = make(chan error, 1)
	t.finish = make(chan bool, 1)
	t.reply = make(chan Reply, 1)
	t.ap = make(chan []RowData, 1)
	t.ctx, t.cancel = context.WithCancel(ctx)
	t.ti = ti

	if t.HasFp() {
		t.fp.sp = int64(math.Ceil(float64(ti.GetTotal()) / t.fp.spn))
	}

}

// GetBaseTask 获取基础任务结构
func (t *BaseTask) GetBaseTask() *BaseTask {
	return t
}

// 监听加入的数据
func (t *BaseTask) listenAppendData() {
	for {
		select {
		case data, ok := <-t.ap:
			if !ok {
				return
			}

			err := t.ti.WriteData(data)
			if err != nil {
				t.err <- err
				return
			}
		case <-t.ctx.Done():
			return
		}
	}
}

// CheckFinish 查看任务有没有完成
func (t *BaseTask) CheckFinish() {
	// 写入总数已经达到总条数，关闭文件和删除任务
	total := t.ti.GetTotal()
	if t.fr >= total || t.fr >= MaxRow {
		t.finish <- true
	} else if t.HasFp() && t.fr%t.fp.sp == 0 { //达到完成部分导出
		t.notifySp()
	}
}

// 获取步进通知
func (t *BaseTask) listenNotifySp() {
	if t.fp == nil {
		return
	}

	for {
		select {
		case nfp, ok := <-t.fp.notifySp:
			if !ok {
				return
			}
			t.fp.fn(t.fp.spn, nfp)
		case <-t.ctx.Done():
			return
		}
	}
}

// 需要通知部分完成
func (t *BaseTask) notifySp() {
	t.fp.l.Lock()
	defer t.fp.l.Unlock()
	t.fp.nfp++
	t.fp.notifySp <- t.fp.nfp
}

// AppendData 追加任务导出数据
func (t *BaseTask) AppendData(data []RowData) {
	t.ap <- data
}

// 任务结束后
func (t *BaseTask) taskEnd(err error) {
	//结束其他进行中的协程任务
	t.cancel()

	//关闭文件资源占用
	t.ti.FileClose()

	//通知结构体
	reply := Reply{
		Filename: t.ti.GetFilename(),
		FilePath: t.ti.GetFullPath(),
		Fr:       t.GetFr(),
	}

	//有异常的时候处理
	if err != nil {
		//删除文件，防止服务器积压
		_ = os.Remove(t.ti.GetFullPath())
		reply.Error = err
	}

	//通知任务回调
	t.reply <- reply
}

// CanWriteRow 是否可以写入数据
func (t *BaseTask) CanWriteRow() bool {
	if t.pe == nil {
		return true
	}

	t.pe.l.Lock()
	defer t.pe.l.Unlock()

	//先处理一下要过滤的数据
	if t.pe.filterNum != 0 && t.pe.filterNum >= t.pe.nft {
		t.fr--
		t.pe.nft++
		return false
	}

	//看看不是已经到要写入的那一行
	if t.pe.begin <= t.fr {
		return true
	}

	return false
}

// GetReply 获取任务结果
func (t *BaseTask) GetReply() <-chan Reply {
	return t.reply
}

// HasFp 是否存在进度推送
func (t *BaseTask) HasFp() bool {
	return t.fp != nil
}

// FrAdd 写入了一条数据
func (t *BaseTask) FrAdd() {
	t.fr++
}

// GetFr 获取写入文件的行数
func (t *BaseTask) GetFr() int64 {
	return t.fr
}

// ManualEnd 手动结束，有错误传入的时候是
func (t *BaseTask) ManualEnd(err error) {
	if err != nil {
		t.err <- err
		return
	}
	t.finish <- true
}

// Release 任务结束后释放
func (t *BaseTask) Release() {
	close(t.reply)
}