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) }