Joey Armstrong | e8c091f | 2023-01-17 16:56:26 -0500 | [diff] [blame] | 1 | package redis |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "sync" |
| 6 | ) |
| 7 | |
| 8 | // ScanIterator is used to incrementally iterate over a collection of elements. |
| 9 | // It's safe for concurrent use by multiple goroutines. |
| 10 | type ScanIterator struct { |
| 11 | mu sync.Mutex // protects Scanner and pos |
| 12 | cmd *ScanCmd |
| 13 | pos int |
| 14 | } |
| 15 | |
| 16 | // Err returns the last iterator error, if any. |
| 17 | func (it *ScanIterator) Err() error { |
| 18 | it.mu.Lock() |
| 19 | err := it.cmd.Err() |
| 20 | it.mu.Unlock() |
| 21 | return err |
| 22 | } |
| 23 | |
| 24 | // Next advances the cursor and returns true if more values can be read. |
| 25 | func (it *ScanIterator) Next(ctx context.Context) bool { |
| 26 | it.mu.Lock() |
| 27 | defer it.mu.Unlock() |
| 28 | |
| 29 | // Instantly return on errors. |
| 30 | if it.cmd.Err() != nil { |
| 31 | return false |
| 32 | } |
| 33 | |
| 34 | // Advance cursor, check if we are still within range. |
| 35 | if it.pos < len(it.cmd.page) { |
| 36 | it.pos++ |
| 37 | return true |
| 38 | } |
| 39 | |
| 40 | for { |
| 41 | // Return if there is no more data to fetch. |
| 42 | if it.cmd.cursor == 0 { |
| 43 | return false |
| 44 | } |
| 45 | |
| 46 | // Fetch next page. |
| 47 | switch it.cmd.args[0] { |
| 48 | case "scan", "qscan": |
| 49 | it.cmd.args[1] = it.cmd.cursor |
| 50 | default: |
| 51 | it.cmd.args[2] = it.cmd.cursor |
| 52 | } |
| 53 | |
| 54 | err := it.cmd.process(ctx, it.cmd) |
| 55 | if err != nil { |
| 56 | return false |
| 57 | } |
| 58 | |
| 59 | it.pos = 1 |
| 60 | |
| 61 | // Redis can occasionally return empty page. |
| 62 | if len(it.cmd.page) > 0 { |
| 63 | return true |
| 64 | } |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | // Val returns the key/field at the current cursor position. |
| 69 | func (it *ScanIterator) Val() string { |
| 70 | var v string |
| 71 | it.mu.Lock() |
| 72 | if it.cmd.Err() == nil && it.pos > 0 && it.pos <= len(it.cmd.page) { |
| 73 | v = it.cmd.page[it.pos-1] |
| 74 | } |
| 75 | it.mu.Unlock() |
| 76 | return v |
| 77 | } |