Golang Iterator Channel Implementation

這是一個關於 Go 的程式研究紀錄筆記，原先設計一個迭代器的操作模式 (Iterator pattern)：

type IntIterator interface {
   Iterator() <-chan int
}

建立使用數字陣列實作 Iterator 來作為例子：

type intIterator []int

func (i intIterator) Iterator() <-chan int {
   c := make(chan int)
   go func() {
      defer close(c)
      for _, v := range i {
         c <- v
      }
   }()
   return c
}

因此可以使用寫下類似 Golish 風格的程式碼，取得並列出所有 Iterator 的數值：

nums := intIterator([]int{1, 2, 3, 5, 7, 11, 13, 17, 19, 23})
for n := range nums.Iterator() {
   println(n)
}

但如果在這種使用情況下，會發現有些東西會不太對：

nums := intIterator([]int{1, 2, 3, 5, 7, 11, 13, 17, 19, 23})
for n := range nums.Iterator() {
   println(n)
   if n == 7 {
      break
   }
}

由於中途離開迴圈，沒有任何的程序會使用或是關閉該 channel c，會造成該 goroutine 會持續停在那裡，造成 goroutine leak，浪費資源。

為了處理這種因為使用來造成的 leak，因此需要借用 context 來得知是否要釋放，修改 Iterator pattern：

import (
   "context"
)

type IntIterator interface {
   Iterator(ctx context.Context) <-chan int
}

調整原先的實作，處理 ctx.Done 的發生：

func (i intIterator) Iterator(ctx context.Context) <-chan int {
   c := make(chan int)
   go func() {
      defer close(c)
      for _, v := range i {
         select {
         case c <- v:
            break
         case <-ctx.Done():
            return
         }
      }
   }()
   return c
}

因此在原先的使用上，建立 context 並傳入給 Iterator，如果就算中途離開 Iterator，其 channel 會能夠被關閉，使得 goroutine leak 消失：

nums := intIterator([]int{1, 2, 3, 5, 7, 11, 13, 17, 19, 23})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
for n := range nums.Iterator(ctx) {
   println(n)
   if n == 7 {
      break
   }
}