package progress import ( "slices" "sync" "time" ) type rawProgressWriter interface { WriteRawProgress(*Progress) error Close() error } type MultiWriter struct { mu sync.Mutex items []*Progress writers map[rawProgressWriter]struct{} meta map[string]any } var _ rawProgressWriter = &MultiWriter{} func NewMultiWriter(opts ...WriterOption) *MultiWriter { mw := &MultiWriter{ writers: map[rawProgressWriter]struct{}{}, meta: map[string]any{}, } for _, o := range opts { o(mw) } return mw } func (ps *MultiWriter) Add(pw Writer) { rw, ok := pw.(rawProgressWriter) if !ok { return } if pws, ok := rw.(*MultiWriter); ok { if pws.contains(ps) { // this would cause a deadlock, so we should panic instead // NOTE: this can be caused by a cycle in the scheduler states, // which is created by a series of unfortunate edge merges panic("multiwriter loop detected") } } ps.mu.Lock() plist := make([]*Progress, 0, len(ps.items)) plist = append(plist, ps.items...) slices.SortFunc(plist, func(a, b *Progress) int { return a.Timestamp.Compare(b.Timestamp) }) for _, p := range plist { rw.WriteRawProgress(p) } ps.writers[rw] = struct{}{} ps.mu.Unlock() } func (ps *MultiWriter) Delete(pw Writer) { rw, ok := pw.(rawProgressWriter) if !ok { return } ps.mu.Lock() delete(ps.writers, rw) ps.mu.Unlock() } func (ps *MultiWriter) Write(id string, v any) error { p := &Progress{ ID: id, Timestamp: time.Now(), Sys: v, meta: ps.meta, } return ps.writeRawProgress(p) } func (ps *MultiWriter) WriteRawProgress(p *Progress) error { return ps.writeRawProgress(p.Decorate(ps.meta)) } func (ps *MultiWriter) writeRawProgress(p *Progress) error { ps.mu.Lock() defer ps.mu.Unlock() ps.items = append(ps.items, p) for w := range ps.writers { if err := w.WriteRawProgress(p); err != nil { return err } } return nil } func (ps *MultiWriter) Close() error { return nil } func (ps *MultiWriter) contains(pw rawProgressWriter) bool { ps.mu.Lock() defer ps.mu.Unlock() _, ok := ps.writers[pw] if ok { return true } for w := range ps.writers { w, ok := w.(*MultiWriter) if !ok { continue } if w.contains(pw) { return true } } return false }