21流机制与状态管理

// StreamReader is the consumer side of an Eino stream.
//
// A StreamReader is read-once: only one goroutine should call Recv, and the
// reader must be closed exactly once (whether the loop finishes normally or
// exits early via break or return).
//
// Typical usage:
//
//  defer sr.Close() // always close, even after io.EOF
//  for {
//      chunk, err := sr.Recv()
//      if errors.Is(err, io.EOF) {
//          break
//      }
//      if err != nil {
//          return err
//      }
//      process(chunk)
//  }
//
// To fan-out a single stream to N independent consumers, call [StreamReader.Copy]
// before any Recv; the original reader becomes unusable after the call.
//
// StreamReaders are created by [Pipe], [StreamReaderFromArray],
// [MergeStreamReaders], [MergeNamedStreamReaders], and [StreamReaderWithConvert].
type StreamReader[T any] struct {
  typ readerType

  st *stream[T]

  ar *arrayReader[T]

  msr *multiStreamReader[T]

  srw *streamReaderWithConvert[T]

  csr *childStreamReader[T]
}

// Pipe creates a new stream with the given capacity that represented with StreamWriter and StreamReader.
// The capacity is the maximum number of items that can be buffered in the stream.
// e.g.
//
//  sr, sw := schema.Pipe[string](3)
//  go func() { // send data
//    defer sw.Close()
//    for i := 0; i < 10; i++ {
//      sw.Send(i, nil)
//    }
//  }
//
//  defer sr.Close()
//  for {
//    chunk, err := sr.Recv()
//    if errors.Is(err, io.EOF) {
//      break
//    }
//    if err != nil {
//      panic(err)
//    }
//    fmt.Println(chunk)
//  }
func Pipe[T any](cap int) (*StreamReader[T], *StreamWriter[T]) {
  stm := newStream[T](cap)
  return stm.asReader(), &StreamWriter[T]{stm: stm}
}

// Copy creates n independent StreamReaders that each receive every element of
// the original stream. The original StreamReader becomes unusable after Copy.
//
// Use Copy when two or more pipeline branches need the same stream —
// for example, when a stream must be fed to both a callback handler and the
// next node in a graph:
//
//  copies := sr.Copy(2)
//  sr1, sr2 := copies[0], copies[1]
//  defer sr1.Close()
//  defer sr2.Close()
//
//  // sr1 and sr2 independently read the same elements
//
// n must be at least 1. If n < 2, the original reader is returned unchanged.
func (sr *StreamReader[T]) Copy(n int) []*StreamReader[T] {
  if n < 2 {
    return []*StreamReader[T]{sr}
  }

  if sr.typ == readerTypeArray {
    ret := make([]*StreamReader[T], n)
    for i, ar := range sr.ar.copy(n) {
      ret[i] = &StreamReader[T]{typ: readerTypeArray, ar: ar}
    }
    return ret
  }

  return copyStreamReaders[T](sr, n)
}

type cpStreamElement[T any] struct {
    once sync.Once            // 保证这一格只从源头读一次
    next *cpStreamElement[T]  // 指向下一格(读到才创建)
    item streamItem[T]        // 缓存下来的那一块数据
}

type parentStreamReader[T any] struct {
  // sr is the original StreamReader.
  sr *StreamReader[T]

  // subStreamList maps each child's index to its latest read chunk.
  // Each value comes from a hidden linked list of cpStreamElement.
  subStreamList []*cpStreamElement[T]

  // closedNum is the count of closed children.
  closedNum uint32
}

type childStreamReader[T any] struct {
  parent *parentStreamReader[T]
  index  int
}

elem.once.Do(func() {
    t, err = p.sr.Recv()          // 只有第一个到这格的 child 真的去源头读
    elem.item = streamItem[T]{chunk: t, err: err}
    if err != io.EOF {
        elem.next = &cpStreamElement[T]{}  // 顺手把下一格铺好
        p.subStreamList[idx] = elem.next
    }
})
// 后到的 child 直接读缓存好的 elem.item,不碰源头

func (p *parentStreamReader[T]) close(idx int) {
  if p.subStreamList[idx] == nil {
    return // avoid close multiple times
  }

  p.subStreamList[idx] = nil

  curClosedNum := atomic.AddUint32(&p.closedNum, 1)

  allClosed := int(curClosedNum) == len(p.subStreamList)
  if allClosed {
    p.sr.Close()
  }
}

// MergeStreamReaders fans in multiple StreamReaders into a single StreamReader.
// Elements from all source streams are interleaved in arrival order (non-deterministic).
// The merged reader reaches EOF only after every source stream has been exhausted.
//
// Callers must still close the merged reader; it propagates the close signal
// to all underlying sources.
//
// Use [MergeNamedStreamReaders] instead when you need to know which source
// stream ended first (it emits a [SourceEOF] per-source EOF rather than
// silently discarding them).
//
// Returns nil if srs is empty.
func MergeStreamReaders[T any](srs []*StreamReader[T]) *StreamReader[T] {
  if len(srs) < 1 {
    return nil
  }

  if len(srs) < 2 {
    return srs[0]
  }

  var arr []T
  var ss []*stream[T]

  for _, sr := range srs {
    switch sr.typ {
    case readerTypeStream:
      ss = append(ss, sr.st)
    case readerTypeArray:
      arr = append(arr, sr.ar.arr[sr.ar.index:]...)
    case readerTypeMultiStream:
      ss = append(ss, sr.msr.nonClosedStreams()...)
    case readerTypeWithConvert:
      ss = append(ss, sr.srw.toStream())
    case readerTypeChild:
      ss = append(ss, sr.csr.toStream())
    default:
      panic("impossible")
    }
  }

  if len(ss) == 0 {
    return &StreamReader[T]{
      typ: readerTypeArray,
      ar: &arrayReader[T]{
        arr:   arr,
        index: 0,
      },
// … 省略 13 行；完整声明 L900–960，点击上方「浏览完整文件」

func (msr *multiStreamReader[T]) recv() (T, error) {
  for len(msr.nonClosed) > 0 {
    var chosen int
    var ok bool
    if len(msr.nonClosed) > maxSelectNum {
      var recv reflect.Value
      chosen, recv, ok = reflect.Select(msr.itemsCases)
      if ok {
        item := recv.Interface().(streamItem[T])
        return item.chunk, item.err
      }
      msr.itemsCases[chosen].Chan = reflect.Value{}
    } else {
      var item *streamItem[T]
      chosen, item, ok = receiveN(msr.nonClosed, msr.sts)
      if ok {
        return item.chunk, item.err
      }
    }

    // delete the closed stream
    for i := range msr.nonClosed {
      if msr.nonClosed[i] == chosen {
        msr.nonClosed = append(msr.nonClosed[:i], msr.nonClosed[i+1:]...)
        break
      }
    }

    if len(msr.sourceReaderNames) > 0 {
      var t T
      return t, &SourceEOF{msr.sourceReaderNames[chosen]}
    }
  }

  var t T
  return t, io.EOF
}

func receiveN[T any](chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
  return []func(chosenList []int, ss []*stream[T]) (index int, item *streamItem[T], ok bool){
    nil,
    func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
      item, ok := <-ss[chosenList[0]].items
      return chosenList[0], &item, ok
    },
    func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
      select {
      case item, ok := <-ss[chosenList[0]].items:
        return chosenList[0], &item, ok
      case item, ok := <-ss[chosenList[1]].items:
        return chosenList[1], &item, ok
      }
    },
    func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
      select {
      case item, ok := <-ss[chosenList[0]].items:
        return chosenList[0], &item, ok
      case item, ok := <-ss[chosenList[1]].items:
        return chosenList[1], &item, ok
      case item, ok := <-ss[chosenList[2]].items:
        return chosenList[2], &item, ok
      }
    },
    func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
      select {
      case item, ok := <-ss[chosenList[0]].items:
        return chosenList[0], &item, ok
      case item, ok := <-ss[chosenList[1]].items:
        return chosenList[1], &item, ok
      case item, ok := <-ss[chosenList[2]].items:
        return chosenList[2], &item, ok
      case item, ok := <-ss[chosenList[3]].items:
        return chosenList[3], &item, ok
      }
    },
    func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) {
      select {
      case item, ok := <-ss[chosenList[0]].items:
        return chosenList[0], &item, ok
      case item, ok := <-ss[chosenList[1]].items:
        return chosenList[1], &item, ok
      case item, ok := <-ss[chosenList[2]].items:
        return chosenList[2], &item, ok
      case item, ok := <-ss[chosenList[3]].items:
        return chosenList[3], &item, ok
      case item, ok := <-ss[chosenList[4]].items:
// … 省略 5 行；完整声明 L21–73，点击上方「浏览完整文件」

const maxSelectNum = 5

// MergeNamedStreamReaders merges multiple named StreamReaders into one.
// Unlike [MergeStreamReaders], when a source stream reaches EOF the merged
// reader emits a [SourceEOF] error (instead of silently continuing) so you can
// detect exactly which source finished. Use [GetSourceName] to retrieve the
// name from a SourceEOF error. The merged reader itself signals io.EOF only
// after all named sources are exhausted.
//
// This is useful when downstream logic must react differently to each source
// completing — for example, draining one agent's output before proceeding:
//
//  namedStreams := map[string]*schema.StreamReader[string]{
//      "agent_a": srA,
//      "agent_b": srB,
//  }
//  merged := schema.MergeNamedStreamReaders(namedStreams)
//  defer merged.Close()
//  for {
//      chunk, err := merged.Recv()
//      if errors.Is(err, io.EOF) { break }
//      if name, ok := schema.GetSourceName(err); ok {
//          fmt.Printf("%s finished\n", name)
//          continue
//      }
//      if err != nil { return err }
//      process(chunk)
//  }
//
// Returns nil if srs is empty.
func MergeNamedStreamReaders[T any](srs map[string]*StreamReader[T]) *StreamReader[T] {
  if len(srs) < 1 {
    return nil
  }

  ss := make([]*StreamReader[T], len(srs))
  names := make([]string, len(srs))

  i := 0
  for name, sr := range srs {
    ss[i] = sr
    names[i] = name
    i++
  }

  return InternalMergeNamedStreamReaders(ss, names)
}

// InternalMergeNamedStreamReaders merges multiple readers with their names
// into a single multi-stream reader.
func InternalMergeNamedStreamReaders[T any](srs []*StreamReader[T], names []string) *StreamReader[T] {
  ss := make([]*stream[T], len(srs))

  for i, sr := range srs {
    ss[i] = sr.toStream()
  }

  msr := newMultiStreamReader(ss)
  msr.sourceReaderNames = names

  return &StreamReader[T]{
    typ: readerTypeMultiStream,
    msr: msr,
  }
}

// StreamReaderFromArray creates a StreamReader from a given slice of elements.
// It takes an array of type T and returns a pointer to a StreamReader[T].
// This allows for streaming the elements of the array in a controlled manner.
// eg.
//
//  sr := schema.StreamReaderFromArray([]int{1, 2, 3})
//  defer sr.Close()
//
//  for {
//    chunk, err := sr.Recv()
//    if errors.Is(err, io.EOF) {
//      break
//    }
//    if err != nil {
//      panic(err)
//    }
//    fmt.Println(chunk)
//  }
func StreamReaderFromArray[T any](arr []T) *StreamReader[T] {
  return &StreamReader[T]{ar: &arrayReader[T]{arr: arr}, typ: readerTypeArray}
}

func concatStreamReader[T any](sr *schema.StreamReader[T]) (T, error) {
  defer sr.Close()

  var items []T

  for {
    chunk, err := sr.Recv()
    if err != nil {
      if err == io.EOF {
        break
      }

      if _, ok := schema.GetSourceName(err); ok {
        continue
      }

      var t T
      return t, newStreamReadError(err)
    }

    items = append(items, chunk)
  }

  if len(items) == 0 {
    var t T
    return t, emptyStreamConcatErr
  }

  if len(items) == 1 {
    return items[0], nil
  }

  res, err := internal.ConcatItems(items)
  if err != nil {
    var t T
    return t, err
  }
  return res, nil
}

// ConcatItems the caller should ensure len(items) > 1
func ConcatItems[T any](items []T) (T, error) {
  typ := generic.TypeOf[T]()
  v := reflect.ValueOf(items)

  var cv reflect.Value
  var err error

  // handle map kind
  if typ.Kind() == reflect.Map {
    cv, err = concatMaps(v)
  } else {
    cv, err = ConcatSliceValue(v)
  }

  if err != nil {
    var t T
    return t, err
  }

  return cv.Interface().(T), nil
}

// WithGenLocalState registers a function to generate per-run local state
// that can be shared across nodes in the graph.
func WithGenLocalState[S any](gls GenLocalState[S]) NewGraphOption {
  return func(ngo *newGraphOptions) {
    ngo.withState = func(ctx context.Context) any {
      return gls(ctx)
    }
    ngo.stateType = generic.TypeOf[S]()
  }
}

func (g *graph) compile(ctx context.Context, opt *graphCompileOptions) (*composableRunnable, error) {
  if g.buildError != nil {
    return nil, g.buildError
  }

  // get run type
  runType := runTypePregel
  cb := pregelChannelBuilder
  if isChain(g.cmp) || isWorkflow(g.cmp) {
    if opt != nil && opt.nodeTriggerMode != "" {
      return nil, errors.New(fmt.Sprintf("%s doesn't support node trigger mode option", g.cmp))
    }
  }
  if (opt != nil && opt.nodeTriggerMode == AllPredecessor) || isWorkflow(g.cmp) {
    runType = runTypeDAG
    cb = dagChannelBuilder
  }

  // get eager type
  eager := false
  if isWorkflow(g.cmp) || runType == runTypeDAG {
    eager = true
  }
  if opt != nil && opt.eagerDisabled {
    eager = false
  }

  if len(g.startNodes) == 0 {
    return nil, errors.New("start node not set")
  }
  if len(g.endNodes) == 0 {
    return nil, errors.New("end node not set")
  }

  // toValidateMap isn't empty means there are nodes that cannot infer type
  for _, v := range g.toValidateMap {
    if len(v) > 0 {
      return nil, fmt.Errorf("some node's input or output types cannot be inferred: %v", g.toValidateMap)
    }
  }

  for key := range g.fieldMappingRecords {
    // not allowed to map multiple fields to the same field
    toMap := make(map[string]bool)
    for _, mapping := range g.fieldMappingRecords[key] {
      if _, ok := toMap[mapping.to]; ok {
        return nil, fmt.Errorf("duplicate mapping target field: %s of node[%s]", mapping.to, key)
      }
// … 省略 171 行；完整声明 L674–892，点击上方「浏览完整文件」

type internalState struct {
    state  any
    mu     sync.Mutex
    parent *internalState   // 指向外层图的 state
}

// ProcessState processes the state from the context in a concurrency-safe way.
// This is the recommended way to access and modify state in custom nodes.
// The provided function handler will be executed with exclusive access to the state (protected by mutex).
//
// State Lookup Behavior:
// - If the requested state type exists in the current graph, it will be returned
// - If not found in current graph, ProcessState will search in parent graph states (for nested graphs)
// - This enables nested graphs to access state from their parent graphs
// - Follows lexical scoping: inner state of the same type shadows outer state
//
// Concurrency Safety:
// - ProcessState automatically locks the mutex of the state being accessed (current or parent level)
// - Each state level has its own mutex, allowing concurrent access to different levels
// - The lock is held for the entire duration of the handler function
//
// Note: This method will report an error if the state type doesn't match or state is not found in the context chain.
//
// Example - Basic usage in a single graph:
//
//  lambdaFunc := func(ctx context.Context, in string, opts ...any) (string, error) {
//    err := compose.ProcessState[*MyState](ctx, func(ctx context.Context, state *MyState) error {
//      // Safely modify state
//      state.Count++
//      return nil
//    })
//    if err != nil {
//      return "", err
//    }
//    return in, nil
//  }
//
// Example - Nested graph accessing parent state:
//
//  // In an inner graph node
//  innerNode := func(ctx context.Context, input string) (string, error) {
//    // Access parent graph's state
//    err := compose.ProcessState[*OuterState](ctx, func(ctx context.Context, s *OuterState) error {
//      s.Counter++  // Safely modify parent state
//      return nil
//    })
//    if err != nil {
//      return "", err
//    }
//
//    // Also access inner graph's own state
//    err = compose.ProcessState[*InnerState](ctx, func(ctx context.Context, s *InnerState) error {
//      s.Data = "processed"
//      return nil
// … 省略 12 行；完整声明 L114–173，点击上方「浏览完整文件」

for interState != nil {
    if cState, ok := interState.state.(S); ok {
        return cState, &interState.mu, nil   // 找到类型匹配的那一层
    }
    interState = interState.parent            // 没匹配,上溯外层
}

type internalState struct {
  state  any
  mu     sync.Mutex
  parent *internalState
}