Skip to content

Instantly share code, notes, and snippets.

@ydnar
Last active April 26, 2023 00:54
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ydnar/b15dd28a3994fb00c3339543c02ca444 to your computer and use it in GitHub Desktop.
Save ydnar/b15dd28a3994fb00c3339543c02ca444 to your computer and use it in GitHub Desktop.
Weighted queue using multi-channel select in Go
package priority
import (
"context"
"sync/atomic"
)
type Priority int
const (
Low Priority = iota
Medium
High
priorityCount
)
type Queue struct {
ch [priorityCount]chan any
n atomic.Int64
}
// NewQueue initializes a Queue with capacity cap.
func NewQueue(cap uint32) *Queue {
q := &Queue{}
for i := Priority(0); i < priorityCount; i++ {
q.ch[i] = make(chan any, cap)
}
return q
}
// Push adds v to the queue, with priority p.
// Push will block until it successfully queues v, or ctx is done.
// Returns the context error, if any.
func (q *Queue) Push(ctx context.Context, v any, p Priority) error {
select {
case q.ch[p] <- v:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// Pop returns the highest priority item from q, blocking until
// one is queued, or ctx is done. Pop will attempt to fairly
// provide queued values with lower priorities.
// Returns the context error, if any.
func (q *Queue) Pop(ctx context.Context) (any, error) {
select {
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[High]:
return v, nil
case v := <-q.ch[Medium]:
return v, nil
case v := <-q.ch[Medium]:
return v, nil
case v := <-q.ch[Medium]:
return v, nil
case v := <-q.ch[Low]:
return v, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
@ydnar
Copy link
Author

ydnar commented Apr 26, 2023

I can’t believe this works SMDH

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment