Go 通過 Map/Filter/ForEach 等流式 API 高效處理資料

kevinwan發表於2022-01-04

什麼是流處理

如果有 java 使用經驗的同學一定會對 java8 的 Stream 讚不絕口,極大的提高了們對於集合型別資料的處理能力。

int sum = widgets.stream()
              .filter(w -> w.getColor() == RED)
              .mapToInt(w -> w.getWeight())
              .sum();

Stream 能讓我們支援鏈式呼叫和函式程式設計的風格來實現資料的處理,看起來資料像是在流水線一樣不斷的實時流轉加工,最終被彙總。Stream 的實現思想就是將資料處理流程抽象成了一個資料流,每次加工後返回一個新的流供使用。

Stream 功能定義

動手寫程式碼之前,先想清楚,把需求理清楚是最重要的一步,我們嘗試代入作者的視角來思考整個元件的實現流程。首先把底層實現的邏輯放一下 ,先嚐試從零開始進行功能定義 stream 功能。

Stream 的工作流程其實也屬於生產消費者模型,整個流程跟工廠中的生產流程非常相似,嘗試先定義一下 Stream 的生命週期:

  1. 建立階段/資料獲取(原料)
  2. 加工階段/中間處理(流水線加工)
  3. 彙總階段/終結操作(最終產品)

下面圍繞 stream 的三個生命週期開始定義 API:

建立階段

為了建立出資料流 stream 這一抽象物件,可以理解為構造器。

我們支援三種方式構造 stream,分別是:切片轉換,channel 轉換,函式式轉換。

注意這個階段的方法都是普通的公開方法,並不繫結 Stream 物件。

// 通過可變引數模式建立 stream
func Just(items ...interface{}) Stream

// 通過 channel 建立 stream
func Range(source <-chan interface{}) Stream

// 通過函式建立 stream
func From(generate GenerateFunc) Stream

// 拼接 stream
func Concat(s Stream, others ...Stream) Stream

加工階段

加工階段需要進行的操作往往對應了我們的業務邏輯,比如:轉換,過濾,去重,排序等等。

這個階段的 API 屬於 method 需要繫結到 Stream 物件上。

結合常用的業務場景進行如下定義:

// 去除重複item
Distinct(keyFunc KeyFunc) Stream
// 按條件過濾item
Filter(filterFunc FilterFunc, opts ...Option) Stream
// 分組
Group(fn KeyFunc) Stream
// 返回前n個元素
Head(n int64) Stream
// 返回後n個元素
Tail(n int64) Stream
// 轉換物件
Map(fn MapFunc, opts ...Option) Stream
// 合併item到slice生成新的stream
Merge() Stream
// 反轉
Reverse() Stream
// 排序
Sort(fn LessFunc) Stream
// 作用在每個item上
Walk(fn WalkFunc, opts ...Option) Stream
// 聚合其他Stream
Concat(streams ...Stream) Stream

加工階段的處理邏輯都會返回一個新的 Stream 物件,這裡有個基本的實現正規化

彙總階段

彙總階段其實就是我們想要的處理結果,比如:是否匹配,統計數量,遍歷等等。

// 檢查是否全部匹配
AllMatch(fn PredicateFunc) bool
// 檢查是否存在至少一項匹配
AnyMatch(fn PredicateFunc) bool
// 檢查全部不匹配
NoneMatch(fn PredicateFunc) bool
// 統計數量
Count() int
// 清空stream
Done()
// 對所有元素執行操作
ForAll(fn ForAllFunc)
// 對每個元素執行操作
ForEach(fn ForEachFunc)

梳理完元件的需求邊界後,我們對於即將要實現的 Stream 有了更清晰的認識。在我的認知裡面真正的架構師對於需求的把握以及後續演化能達到及其精準的地步,做到這一點離不開對需求的深入思考以及洞穿需求背後的本質。通過代入作者的視角來模擬覆盤整個專案的構建流程,學習作者的思維方法論這正是我們學習開源專案最大的價值所在。

好了,我們嘗試定義出完整的 Stream 介面全貌以及函式。

介面的作用不僅僅是模版作用,還在於利用其抽象能力搭建專案整體的框架而不至於一開始就陷入細節,能快速的將我們的思考過程通過介面簡潔的表達出來,學會養成自頂向下的思維方法從巨集觀的角度來觀察整個系統,一開始就陷入細節則很容易拔劍四顧心茫然。。。
rxOptions struct {
  unlimitedWorkers bool
  workers          int
}
Option func(opts *rxOptions)
// key生成器
//item - stream中的元素
KeyFunc func(item interface{}) interface{}
// 過濾函式
FilterFunc func(item interface{}) bool
// 物件轉換函式
MapFunc func(intem interface{}) interface{}
// 物件比較
LessFunc func(a, b interface{}) bool
// 遍歷函式
WalkFunc func(item interface{}, pip chan<- interface{})
// 匹配函式
PredicateFunc func(item interface{}) bool
// 對所有元素執行操作
ForAllFunc func(pip <-chan interface{})
// 對每個item執行操作
ForEachFunc func(item interface{})
// 對每個元素併發執行操作
ParallelFunc func(item interface{})
// 對所有元素執行聚合操作
ReduceFunc func(pip <-chan interface{}) (interface{}, error)
// item生成函式
GenerateFunc func(source <-chan interface{})

Stream interface {
  // 去除重複item
  Distinct(keyFunc KeyFunc) Stream
  // 按條件過濾item
  Filter(filterFunc FilterFunc, opts ...Option) Stream
  // 分組
  Group(fn KeyFunc) Stream
  // 返回前n個元素
  Head(n int64) Stream
  // 返回後n個元素
  Tail(n int64) Stream
  // 獲取第一個元素
  First() interface{}
  // 獲取最後一個元素
  Last() interface{}
  // 轉換物件
  Map(fn MapFunc, opts ...Option) Stream
  // 合併item到slice生成新的stream
  Merge() Stream
  // 反轉
  Reverse() Stream
  // 排序
  Sort(fn LessFunc) Stream
  // 作用在每個item上
  Walk(fn WalkFunc, opts ...Option) Stream
  // 聚合其他Stream
  Concat(streams ...Stream) Stream
  // 檢查是否全部匹配
  AllMatch(fn PredicateFunc) bool
  // 檢查是否存在至少一項匹配
  AnyMatch(fn PredicateFunc) bool
  // 檢查全部不匹配
  NoneMatch(fn PredicateFunc) bool
  // 統計數量
  Count() int
  // 清空stream
  Done()
  // 對所有元素執行操作
  ForAll(fn ForAllFunc)
  // 對每個元素執行操作
  ForEach(fn ForEachFunc)
}

channel() 方法用於獲取 Stream 管道屬性,因為在具體實現時我們面向的是介面物件所以暴露一個私有方法 read 出來。

// 獲取內部的資料容器channel,內部方法
channel() chan interface{}

實現思路

功能定義梳理清楚了,接下來考慮幾個工程實現的問題。

如何實現鏈式呼叫

鏈式呼叫,建立物件用到的 builder 模式可以達到鏈式呼叫效果。實際上 Stream 實現類似鏈式的效果原理也是一樣的,每次呼叫完後都建立一個新的 Stream 返回給使用者。

// 去除重複item
Distinct(keyFunc KeyFunc) Stream
// 按條件過濾item
Filter(filterFunc FilterFunc, opts ...Option) Stream

如何實現流水線的處理效果

所謂的流水線可以理解為資料在 Stream 中的儲存容器,在 go 中我們可以使用 channel 作為資料的管道,達到 Stream 鏈式呼叫執行多個操作時非同步非阻塞效果。

如何支援並行處理

資料加工本質上是在處理 channel 中的資料,那麼要實現並行處理無非是並行消費 channel 而已,利用 goroutine 協程、WaitGroup 機制可以非常方便的實現並行處理。

go-zero 實現

core/fx/stream.go

go-zero 中關於 Stream 的實現並沒有定義介面,不過沒關係底層實現時邏輯是一樣的。

為了實現 Stream 介面我們定義一個內部的實現類,其中 source 為 channel 型別,模擬流水線功能。

Stream struct {
  source <-chan interface{}
}

建立 API

channel 建立 Range

通過 channel 建立 stream

func Range(source <-chan interface{}) Stream {  
  return Stream{  
    source: source,  
  }  
}

可變引數模式建立 Just

通過可變引數模式建立 stream,channel 寫完後及時 close 是個好習慣。

func Just(items ...interface{}) Stream {
  source := make(chan interface{}, len(items))
  for _, item := range items {
    source <- item
  }
  close(source)
  return Range(source)
}

函式建立 From

通過函式建立 Stream

func From(generate GenerateFunc) Stream {
  source := make(chan interface{})
  threading.GoSafe(func() {
    defer close(source)
    generate(source)
  })
  return Range(source)
}

因為涉及外部傳入的函式引數呼叫,執行過程並不可用因此需要捕捉執行時異常防止 panic 錯誤傳導到上層導致應用崩潰。

func Recover(cleanups ...func()) {
  for _, cleanup := range cleanups {
    cleanup()
  }
  if r := recover(); r != nil {
    logx.ErrorStack(r)
  }
}

func RunSafe(fn func()) {
  defer rescue.Recover()
  fn()
}

func GoSafe(fn func()) {
  go RunSafe(fn)
}

拼接 Concat

拼接其他 Stream 建立一個新的 Stream,呼叫內部 Concat method 方法,後文將會分析 Concat 的原始碼實現。

func Concat(s Stream, others ...Stream) Stream {
  return s.Concat(others...)
}

加工 API

去重 Distinct

因為傳入的是函式引數KeyFunc func(item interface{}) interface{}意味著也同時支援按照業務場景自定義去重,本質上是利用 KeyFunc 返回的結果基於 map 實現去重。

函式引數非常強大,能極大的提升靈活性。

func (s Stream) Distinct(keyFunc KeyFunc) Stream {
  source := make(chan interface{})
  threading.GoSafe(func() {
    // channel記得關閉是個好習慣
    defer close(source)
    keys := make(map[interface{}]lang.PlaceholderType)
    for item := range s.source {
      // 自定義去重邏輯
      key := keyFunc(item)
      // 如果key不存在,則將資料寫入新的channel
      if _, ok := keys[key]; !ok {
        source <- item
        keys[key] = lang.Placeholder
      }
    }
  })
  return Range(source)
}

使用案例:

// 1 2 3 4 5
Just(1, 2, 3, 3, 4, 5, 5).Distinct(func(item interface{}) interface{} {
  return item
}).ForEach(func(item interface{}) {
  t.Log(item)
})

// 1 2 3 4
Just(1, 2, 3, 3, 4, 5, 5).Distinct(func(item interface{}) interface{} {
  uid := item.(int)
  // 對大於4的item進行特殊去重邏輯,最終只保留一個>3的item
  if uid > 3 {
    return 4
  }
  return item
}).ForEach(func(item interface{}) {
  t.Log(item)
})

過濾 Filter

通過將過濾邏輯抽象成 FilterFunc,然後分別作用在 item 上根據 FilterFunc 返回的布林值決定是否寫回新的 channel 中實現過濾功能,實際的過濾邏輯委託給了 Walk method。

Option 引數包含兩個選項:

  1. unlimitedWorkers 不限制協程數量
  2. workers 限制協程數量
FilterFunc func(item interface{}) bool

func (s Stream) Filter(filterFunc FilterFunc, opts ...Option) Stream {
  return s.Walk(func(item interface{}, pip chan<- interface{}) {
    if filterFunc(item) {
      pip <- item
    }
  }, opts...)
}

使用示例:

func TestInternalStream_Filter(t *testing.T) {
  // 保留偶數 2,4
  channel := Just(1, 2, 3, 4, 5).Filter(func(item interface{}) bool {
    return item.(int)%2 == 0
  }).channel()
  for item := range channel {
    t.Log(item)
  }
}

遍歷執行 Walk

walk 英文意思是步行,這裡的意思是對每個 item 都執行一次 WalkFunc 操作並將結果寫入到新的 Stream 中。

這裡注意一下因為內部採用了協程機制非同步執行讀取和寫入資料所以新的 Stream 中 channel 裡面的資料順序是隨機的。

// item-stream中的item元素
// pipe-item符合條件則寫入pipe
WalkFunc func(item interface{}, pipe chan<- interface{})

func (s Stream) Walk(fn WalkFunc, opts ...Option) Stream {
  option := buildOptions(opts...)
  if option.unlimitedWorkers {
    return s.walkUnLimited(fn, option)
  }
  return s.walkLimited(fn, option)
}

func (s Stream) walkUnLimited(fn WalkFunc, option *rxOptions) Stream {
  // 建立帶緩衝區的channel
  // 預設為16,channel中元素超過16將會被阻塞
  pipe := make(chan interface{}, defaultWorkers)
  go func() {
    var wg sync.WaitGroup

    for item := range s.source {
      // 需要讀取s.source的所有元素
      // 這裡也說明了為什麼channel最後寫完記得完畢
      // 如果不關閉可能導致協程一直阻塞導致洩漏
      // 重要, 不賦值給val是個典型的併發陷阱,後面在另一個goroutine裡使用了
      val := item
      wg.Add(1)
      // 安全模式下執行函式
      threading.GoSafe(func() {
        defer wg.Done()
        fn(item, pipe)
      })
    }
    wg.Wait()
    close(pipe)
  }()

  // 返回新的Stream
  return Range(pipe)
}

func (s Stream) walkLimited(fn WalkFunc, option *rxOptions) Stream {
  pipe := make(chan interface{}, option.workers)
  go func() {
    var wg sync.WaitGroup
    // 控制協程數量
    pool := make(chan lang.PlaceholderType, option.workers)

    for item := range s.source {
      // 重要, 不賦值給val是個典型的併發陷阱,後面在另一個goroutine裡使用了
      val := item
      // 超過協程限制時將會被阻塞
      pool <- lang.Placeholder
      // 這裡也說明了為什麼channel最後寫完記得完畢
      // 如果不關閉可能導致協程一直阻塞導致洩漏
      wg.Add(1)

      // 安全模式下執行函式
      threading.GoSafe(func() {
        defer func() {
          wg.Done()
          //執行完成後讀取一次pool釋放一個協程位置
          <-pool
        }()
        fn(item, pipe)
      })
    }
    wg.Wait()
    close(pipe)
  }()
  return Range(pipe)
}

使用案例:

返回的順序是隨機的。

func Test_Stream_Walk(t *testing.T) {
  // 返回 300,100,200
  Just(1, 2, 3).Walk(func(item interface{}, pip chan<- interface{}) {
    pip <- item.(int) * 100
  }, WithWorkers(3)).ForEach(func(item interface{}) {
    t.Log(item)
  })
}

分組 Group

通過對 item 匹配放入 map 中。

KeyFunc func(item interface{}) interface{}

func (s Stream) Group(fn KeyFunc) Stream {
  groups := make(map[interface{}][]interface{})
  for item := range s.source {
    key := fn(item)
    groups[key] = append(groups[key], item)
  }
  source := make(chan interface{})
  go func() {
    for _, group := range groups {
      source <- group
    }
    close(source)
  }()
  return Range(source)
}

獲取前 n 個元素 Head

n 大於實際資料集長度的話將會返回全部元素

func (s Stream) Head(n int64) Stream {
  if n < 1 {
    panic("n must be greather than 1")
  }
  source := make(chan interface{})
  go func() {
    for item := range s.source {
      n--
      // n值可能大於s.source長度,需要判斷是否>=0
      if n >= 0 {
        source <- item
      }
      // let successive method go ASAP even we have more items to skip
      // why we don't just break the loop, because if break,
      // this former goroutine will block forever, which will cause goroutine leak.
      // n==0說明source已經寫滿可以進行關閉了
      // 既然source已經滿足條件了為什麼不直接進行break跳出迴圈呢?
      // 作者提到了防止協程洩漏
      // 因為每次操作最終都會產生一個新的Stream,舊的Stream永遠也不會被呼叫了
      if n == 0 {
        close(source)
        break
      }
    }
    // 上面的迴圈跳出來了說明n大於s.source實際長度
    // 依舊需要顯示關閉新的source
    if n > 0 {
      close(source)
    }
  }()
  return Range(source)
}

使用示例:

// 返回1,2
func TestInternalStream_Head(t *testing.T) {
  channel := Just(1, 2, 3, 4, 5).Head(2).channel()
  for item := range channel {
    t.Log(item)
  }
}

獲取後 n 個元素 Tail

這裡很有意思,為了確保拿到最後 n 個元素使用環形切片 Ring 這個資料結構,先了解一下 Ring 的實現。

// 環形切片
type Ring struct {
  elements []interface{}
  index    int
  lock     sync.Mutex
}

func NewRing(n int) *Ring {
  if n < 1 {
    panic("n should be greather than 0")
  }
  return &Ring{
    elements: make([]interface{}, n),
  }
}

// 新增元素
func (r *Ring) Add(v interface{}) {
  r.lock.Lock()
  defer r.lock.Unlock()
  // 將元素寫入切片指定位置
  // 這裡的取餘實現了迴圈寫效果
  r.elements[r.index%len(r.elements)] = v
  // 更新下次寫入位置
  r.index++
}

// 獲取全部元素
// 讀取順序保持與寫入順序一致
func (r *Ring) Take() []interface{} {
  r.lock.Lock()
  defer r.lock.Unlock()

  var size int
  var start int
  // 當出現迴圈寫的情況時
  // 開始讀取位置需要通過去餘實現,因為我們希望讀取出來的順序與寫入順序一致
  if r.index > len(r.elements) {
    size = len(r.elements)
    // 因為出現迴圈寫情況,當前寫入位置index開始為最舊的資料
    start = r.index % len(r.elements)
  } else {
    size = r.index
  }
  elements := make([]interface{}, size)
  for i := 0; i < size; i++ {
    // 取餘實現環形讀取,讀取順序保持與寫入順序一致
    elements[i] = r.elements[(start+i)%len(r.elements)]
  }

  return elements
}

總結一下環形切片的優點:

  • 支援自動滾動更新
  • 節省記憶體

環形切片能實現固定容量滿的情況下舊資料不斷被新資料覆蓋,由於這個特性可以用於讀取 channel 後 n 個元素。

func (s Stream) Tail(n int64) Stream {
  if n < 1 {
    panic("n must be greather than 1")
  }
  source := make(chan interface{})
  go func() {
    ring := collection.NewRing(int(n))
    // 讀取全部元素,如果數量>n環形切片能實現新資料覆蓋舊資料
    // 保證獲取到的一定最後n個元素
    for item := range s.source {
      ring.Add(item)
    }
    for _, item := range ring.Take() {
      source <- item
    }
    close(source)
  }()
  return Range(source)
}

那麼為什麼不直接使用 len(source) 長度的切片呢?

答案是節省記憶體。凡是涉及到環形型別的資料結構時都具備一個優點那就省記憶體,能做到按需分配資源。

使用示例:

func TestInternalStream_Tail(t *testing.T) {
  // 4,5
  channel := Just(1, 2, 3, 4, 5).Tail(2).channel()
  for item := range channel {
    t.Log(item)
  }
  // 1,2,3,4,5
  channel2 := Just(1, 2, 3, 4, 5).Tail(6).channel()
  for item := range channel2 {
    t.Log(item)
  }
}

元素轉換Map

元素轉換,內部由協程完成轉換操作,注意輸出channel並不保證按原序輸出。

MapFunc func(intem interface{}) interface{}
func (s Stream) Map(fn MapFunc, opts ...Option) Stream {
  return s.Walk(func(item interface{}, pip chan<- interface{}) {
    pip <- fn(item)
  }, opts...)
}

使用示例:

func TestInternalStream_Map(t *testing.T) {
  channel := Just(1, 2, 3, 4, 5, 2, 2, 2, 2, 2, 2).Map(func(item interface{}) interface{} {
    return item.(int) * 10
  }).channel()
  for item := range channel {
    t.Log(item)
  }
}

合併 Merge

實現比較簡單,我考慮了很久沒想到有什麼場景適合這個方法。

func (s Stream) Merge() Stream {
  var items []interface{}
  for item := range s.source {
    items = append(items, item)
  }
  source := make(chan interface{}, 1)
  source <- items
  return Range(source)
}

反轉 Reverse

反轉 channel 中的元素。反轉演算法流程是:

  • 找到中間節點
  • 節點兩邊開始兩兩交換

注意一下為什麼獲取 s.source 時用切片來接收呢? 切片會自動擴容,用陣列不是更好嗎?

其實這裡是不能用陣列的,因為不知道 Stream 寫入 source 的操作往往是在協程非同步寫入的,每個 Stream 中的 channel 都可能在動態變化,用流水線來比喻 Stream 工作流程的確非常形象。

func (s Stream) Reverse() Stream {
  var items []interface{}
  for item := range s.source {
    items = append(items, item)
  }
  for i := len(items)/2 - 1; i >= 0; i-- {
    opp := len(items) - 1 - i
    items[i], items[opp] = items[opp], items[i]
  }
  return Just(items...)
}

使用示例:

func TestInternalStream_Reverse(t *testing.T) {
  channel := Just(1, 2, 3, 4, 5).Reverse().channel()
  for item := range channel {
    t.Log(item)
  }
}

排序 Sort

內網呼叫 slice 官方包的排序方案,傳入比較函式實現比較邏輯即可。

func (s Stream) Sort(fn LessFunc) Stream {
  var items []interface{}
  for item := range s.source {
    items = append(items, item)
  }

  sort.Slice(items, func(i, j int) bool {
    return fn(i, j)
  })
  return Just(items...)
}

使用示例:

// 5,4,3,2,1
func TestInternalStream_Sort(t *testing.T) {
  channel := Just(1, 2, 3, 4, 5).Sort(func(a, b interface{}) bool {
    return a.(int) > b.(int)
  }).channel()
  for item := range channel {
    t.Log(item)
  }
}

拼接 Concat

func (s Stream) Concat(steams ...Stream) Stream {
  // 建立新的無緩衝channel
  source := make(chan interface{})
  go func() {
    // 建立一個waiGroup物件
    group := threading.NewRoutineGroup()
    // 非同步從原channel讀取資料
    group.Run(func() {
      for item := range s.source {
        source <- item
      }
    })
    // 非同步讀取待拼接Stream的channel資料
    for _, stream := range steams {
      // 每個Stream開啟一個協程
      group.Run(func() {
        for item := range stream.channel() {
          source <- item
        }
      })
    }
    // 阻塞等待讀取完成
    group.Wait()
    close(source)
  }()
  // 返回新的Stream
  return Range(source)
}

彙總 API

全部匹配 AllMatch

func (s Stream) AllMatch(fn PredicateFunc) bool {
  for item := range s.source {
    if !fn(item) {
      // 需要排空 s.source,否則前面的goroutine可能阻塞
      go drain(s.source)
      return false
    }
  }

  return true
}

任意匹配 AnyMatch

func (s Stream) AnyMatch(fn PredicateFunc) bool {
  for item := range s.source {
    if fn(item) {
      // 需要排空 s.source,否則前面的goroutine可能阻塞
      go drain(s.source)
      return true
    }
  }

  return false
}

一個也不匹配 NoneMatch

func (s Stream) NoneMatch(fn func(item interface{}) bool) bool {
  for item := range s.source {
    if fn(item) {
      // 需要排空 s.source,否則前面的goroutine可能阻塞
      go drain(s.source)
      return false
    }
  }

  return true
}

數量統計 Count

func (s Stream) Count() int {
  var count int
  for range s.source {
    count++
  }
  return count
}

清空 Done

func (s Stream) Done() {
  // 排空 channel,防止 goroutine 阻塞洩露
  drain(s.source)
}

迭代全部元素 ForAll

func (s Stream) ForAll(fn ForAllFunc) {
  fn(s.source)
}

迭代每個元素 ForEach

func (s Stream) ForAll(fn ForAllFunc) {
  fn(s.source)
}

小結

至此 Stream 元件就全部實現完了,核心邏輯是利用 channel 當做管道,資料當做水流,不斷的用協程接收/寫入資料到 channel 中達到非同步非阻塞的效果。

回到開篇提到的問題,未動手前想要實現一個 stream 難度似乎非常大,很難想象在 go 中 300 多行的程式碼就能實現如此強大的元件。

實現高效的基礎來源三個語言特性:

  • channel
  • 協程
  • 函數語言程式設計

參考資料

pipeline模式

切片反轉演算法

專案地址

https://github.com/zeromicro/go-zero

歡迎使用 go-zerostar 支援我們!

微信交流群

關注『微服務實踐』公眾號並點選 交流群 獲取社群群二維碼。

相關文章