github.com/alibaba/ilogtail/pkg@v0.0.0-20250526110833-c53b480d046c/helper/collector_imp.go (about)

     1  // Copyright 2024 iLogtail Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package helper
    16  
    17  import (
    18  	"github.com/alibaba/ilogtail/pkg/models"
    19  	"github.com/alibaba/ilogtail/pkg/pipeline"
    20  )
    21  
    22  // Observable pipeline data collector, which stores data based on channal and can be subscribed by multiple consumers
    23  type observePipeCollector struct {
    24  	groupChan chan *models.PipelineGroupEvents
    25  }
    26  
    27  func (p *observePipeCollector) Collect(group *models.GroupInfo, events ...models.PipelineEvent) {
    28  	if len(events) == 0 {
    29  		return
    30  	}
    31  	p.groupChan <- &models.PipelineGroupEvents{
    32  		Group:  group,
    33  		Events: events,
    34  	}
    35  }
    36  
    37  func (p *observePipeCollector) CollectList(groups ...*models.PipelineGroupEvents) {
    38  	if len(groups) == 0 {
    39  		return
    40  	}
    41  	for _, g := range groups {
    42  		p.groupChan <- g
    43  	}
    44  }
    45  
    46  func (p *observePipeCollector) ToArray() []*models.PipelineGroupEvents {
    47  	totalCount := len(p.groupChan)
    48  	results := make([]*models.PipelineGroupEvents, totalCount)
    49  	for i := 0; i < totalCount; i++ {
    50  		results[i] = <-p.groupChan
    51  	}
    52  	return results
    53  }
    54  
    55  func (p *observePipeCollector) Observe() chan *models.PipelineGroupEvents {
    56  	return p.groupChan
    57  }
    58  
    59  func (p *observePipeCollector) Close() {
    60  	close(p.groupChan)
    61  }
    62  
    63  // groupedPipeCollector group the collected PipelineEvent by groupInfo.
    64  // The limitation is that it cannot be subscribed as it always returns an empty chan.
    65  // so it can only return all the grouped data at one time.
    66  type groupedPipeCollector struct {
    67  	groupEvents map[*models.GroupInfo][]models.PipelineEvent
    68  }
    69  
    70  func (p *groupedPipeCollector) Collect(group *models.GroupInfo, events ...models.PipelineEvent) {
    71  	if len(events) == 0 {
    72  		return
    73  	}
    74  	store, has := p.groupEvents[group]
    75  	if !has {
    76  		store = make([]models.PipelineEvent, 0)
    77  	}
    78  	p.groupEvents[group] = append(store, events...)
    79  }
    80  
    81  func (p *groupedPipeCollector) CollectList(groups ...*models.PipelineGroupEvents) {
    82  	if len(groups) == 0 {
    83  		return
    84  	}
    85  	for _, g := range groups {
    86  		p.Collect(g.Group, g.Events...)
    87  	}
    88  }
    89  
    90  func (p *groupedPipeCollector) ToArray() []*models.PipelineGroupEvents {
    91  	len, idx := len(p.groupEvents), 0
    92  	results := make([]*models.PipelineGroupEvents, len)
    93  	if len == 0 {
    94  		return results
    95  	}
    96  	for group, events := range p.groupEvents {
    97  		results[idx] = &models.PipelineGroupEvents{
    98  			Group:  group,
    99  			Events: events,
   100  		}
   101  		idx++
   102  	}
   103  	p.groupEvents = make(map[*models.GroupInfo][]models.PipelineEvent)
   104  	return results
   105  }
   106  
   107  func (p *groupedPipeCollector) Observe() chan *models.PipelineGroupEvents {
   108  	return nil
   109  }
   110  
   111  func (p *groupedPipeCollector) Close() {
   112  	for k := range p.groupEvents {
   113  		delete(p.groupEvents, k)
   114  	}
   115  }
   116  
   117  // noopPipeCollector is an empty collector implementation.
   118  type noopPipeCollector struct {
   119  }
   120  
   121  func (p *noopPipeCollector) Collect(group *models.GroupInfo, events ...models.PipelineEvent) {
   122  }
   123  
   124  func (p *noopPipeCollector) CollectList(groups ...*models.PipelineGroupEvents) {
   125  }
   126  
   127  func (p *noopPipeCollector) ToArray() []*models.PipelineGroupEvents {
   128  	return nil
   129  }
   130  
   131  func (p *noopPipeCollector) Observe() chan *models.PipelineGroupEvents {
   132  	return nil
   133  }
   134  
   135  func (p *noopPipeCollector) Close() {
   136  }
   137  
   138  type defaultPipelineContext struct {
   139  	collector pipeline.PipelineCollector
   140  }
   141  
   142  func (p *defaultPipelineContext) Collector() pipeline.PipelineCollector {
   143  	return p.collector
   144  }
   145  
   146  func NewObservePipelineContext(queueSize int) pipeline.PipelineContext {
   147  	return newPipelineContext(&observePipeCollector{
   148  		groupChan: make(chan *models.PipelineGroupEvents, queueSize),
   149  	})
   150  }
   151  
   152  func NewGroupedPipelineContext() pipeline.PipelineContext {
   153  	return newPipelineContext(&groupedPipeCollector{
   154  		groupEvents: make(map[*models.GroupInfo][]models.PipelineEvent),
   155  	})
   156  }
   157  
   158  func NewNoopPipelineContext() pipeline.PipelineContext {
   159  	return newPipelineContext(&noopPipeCollector{})
   160  }
   161  
   162  func newPipelineContext(collector pipeline.PipelineCollector) pipeline.PipelineContext {
   163  	return &defaultPipelineContext{collector: collector}
   164  }