-
Notifications
You must be signed in to change notification settings - Fork 38
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
5e8e186
commit 3787392
Showing
3 changed files
with
262 additions
and
48 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,198 @@ | ||
| package merkledag | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
| "sync/atomic" | ||
|
|
||
| cid "github.com/ipfs/go-cid" | ||
| ) | ||
|
|
||
| // contextKey is a type to use as value for the ProgressTracker contexts. | ||
| type progressTrackerContextKey string | ||
|
|
||
| const progressContextKey progressTrackerContextKey = "progress" | ||
|
|
||
| // ProgressTracker is used to show progress when fetching nodes. | ||
| type ProgressTracker struct { | ||
| Events chan *ProgressTrackerEvent | ||
|
|
||
| total int32 | ||
|
|
||
| trackEvents bool | ||
| lk sync.Mutex | ||
|
|
||
| bufferedEvents ProgressTrackerEvents | ||
| inEvents chan *ProgressTrackerEvent | ||
| newEventsSignal chan struct{} | ||
|
|
||
| onceStop sync.Once | ||
| stop chan struct{} | ||
| stopped chan struct{} | ||
| } | ||
|
|
||
| // WithProgressTracker returns a new context with value "progress" derived from | ||
| // the given one. | ||
| func WithProgressTracker(ctx context.Context, p *ProgressTracker) (nCtx context.Context, cancel func()) { | ||
| return context.WithValue(ctx, progressContextKey, p), func() { p.Stop() } | ||
| } | ||
|
|
||
| // NewProgressTracker returns progress tracker instance. Verbose turns on events tracking | ||
| func NewProgressTracker(trackEvents bool) *ProgressTracker { | ||
| tracker := &ProgressTracker{ | ||
| Events: make(chan *ProgressTrackerEvent), | ||
| trackEvents: trackEvents, | ||
| inEvents: make(chan *ProgressTrackerEvent), | ||
| newEventsSignal: make(chan struct{}, 1), | ||
|
|
||
| stop: make(chan struct{}), | ||
| stopped: make(chan struct{}, 2), // in, out pumps | ||
| } | ||
|
|
||
| if trackEvents { | ||
| go tracker.inPump() | ||
| go tracker.outPump() | ||
| } | ||
|
|
||
| return tracker | ||
| } | ||
|
|
||
| // GetProgressTracker returns a progress tracker instance if present | ||
| func GetProgressTracker(ctx context.Context) *ProgressTracker { | ||
| v, _ := ctx.Value(progressContextKey).(*ProgressTracker) | ||
| return v | ||
| } | ||
|
|
||
| // Value returns the current progress. | ||
| func (p *ProgressTracker) Value() int { | ||
| return int(atomic.LoadInt32(&p.total)) | ||
| } | ||
|
|
||
| // Stop stops the "pump" goroutines and clear the resources | ||
| func (p *ProgressTracker) Stop() { | ||
| if p.trackEvents { | ||
| return | ||
| } | ||
|
|
||
| p.onceStop.Do(func() { | ||
| close(p.stop) | ||
| }) | ||
|
|
||
| // in, out pumps | ||
| <-p.stopped | ||
| <-p.stopped | ||
|
|
||
| close(p.stopped) | ||
| } | ||
|
|
||
| func (p *ProgressTracker) addNewEventToBuffer(event *ProgressTrackerEvent) { | ||
| p.lk.Lock() | ||
| p.bufferedEvents = append(p.bufferedEvents, event) | ||
| p.lk.Unlock() | ||
|
|
||
| select { | ||
| case p.newEventsSignal <- struct{}{}: | ||
| default: | ||
| } | ||
| } | ||
|
|
||
| func (p *ProgressTracker) inPump() { | ||
| defer func() { | ||
| p.stopped <- struct{}{} | ||
| close(p.newEventsSignal) | ||
| }() | ||
|
|
||
| for { | ||
| select { | ||
| case event, more := <-p.inEvents: | ||
| if !more { | ||
| return | ||
| } | ||
|
|
||
| p.addNewEventToBuffer(event) | ||
| case <-p.stop: | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (p *ProgressTracker) outPump() { | ||
| defer func() { | ||
| p.stopped <- struct{}{} | ||
| close(p.Events) | ||
| }() | ||
|
|
||
| finished := false | ||
| for { | ||
| _, more := <-p.newEventsSignal | ||
| if !more { | ||
| finished = true | ||
| } | ||
|
|
||
| p.lk.Lock() | ||
| eventsToSend := p.bufferedEvents | ||
| p.bufferedEvents = p.bufferedEvents[:0] | ||
| p.lk.Unlock() | ||
|
|
||
| for _, event := range eventsToSend { | ||
| select { | ||
| case p.Events <- event: | ||
| case <-p.stop: | ||
| return | ||
| } | ||
| } | ||
|
|
||
| if finished { | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (p *ProgressTracker) plannedToFetch(c cid.Cid) { | ||
| if !p.trackEvents { | ||
| return | ||
| } | ||
|
|
||
| p.inEvents <- &ProgressTrackerEvent{ | ||
| Type: ProgressTrackerEventTypePlannedToFetch, | ||
| Cid: c, | ||
| } | ||
| } | ||
|
|
||
| func (p *ProgressTracker) fetched(c cid.Cid) { | ||
| atomic.AddInt32(&p.total, 1) | ||
|
|
||
| if !p.trackEvents { | ||
| return | ||
| } | ||
|
|
||
| p.inEvents <- &ProgressTrackerEvent{ | ||
| Type: ProgressTrackerEventTypeFetched, | ||
| Cid: c, | ||
| } | ||
| } | ||
|
|
||
| func (p *ProgressTracker) allFinished() { | ||
| close(p.inEvents) | ||
| } | ||
|
|
||
| const ( | ||
| // ProgressTrackerEventTypePlannedToFetch signals that a cid is planned to fetch | ||
| ProgressTrackerEventTypePlannedToFetch ProgressTrackerEventType = "planned to fetch" | ||
| // ProgressTrackerEventTypeFetched signals that a cid was fetched | ||
| ProgressTrackerEventTypeFetched ProgressTrackerEventType = "was fetched" | ||
| ) | ||
|
|
||
| type ( | ||
| // ProgressTrackerEventType describes an event type | ||
| ProgressTrackerEventType string | ||
|
|
||
| // ProgressTrackerEvent is used to show what exactly has happened | ||
| ProgressTrackerEvent struct { | ||
| Type ProgressTrackerEventType | ||
| Cid cid.Cid | ||
| } | ||
|
|
||
| // ProgressTrackerEvents is a list of events in a chronological order | ||
| ProgressTrackerEvents []*ProgressTrackerEvent | ||
| ) |