github.com/matrixorigin/matrixone@v1.2.0/pkg/util/trace/impl/motrace/Aggr.go (about)

     1  // Copyright 2023 Matrix Origin
     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 motrace
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    22  )
    23  
    24  type Item interface {
    25  	Key(duration time.Duration) interface{}
    26  }
    27  
    28  type Aggregator struct {
    29  	ctx         context.Context
    30  	Grouped     map[interface{}]Item
    31  	WindowSize  time.Duration
    32  	NewItemFunc func(i Item, ctx context.Context) Item
    33  	UpdateFunc  func(ctx context.Context, existing, new Item)
    34  	FilterFunc  func(i Item) bool
    35  }
    36  
    37  type key int
    38  
    39  const (
    40  	DurationKey key = iota
    41  )
    42  
    43  func NewAggregator(ctx context.Context, windowSize time.Duration, newItemFunc func(i Item, ctx context.Context) Item, updateFunc func(ctx context.Context, existing, new Item), filterFunc func(i Item) bool) *Aggregator {
    44  	ctx = context.WithValue(ctx, DurationKey, windowSize)
    45  	return &Aggregator{
    46  		ctx:         ctx,
    47  		Grouped:     make(map[interface{}]Item),
    48  		WindowSize:  windowSize,
    49  		NewItemFunc: newItemFunc,
    50  		UpdateFunc:  updateFunc,
    51  		FilterFunc:  filterFunc,
    52  	}
    53  }
    54  
    55  var ErrFilteredOut = moerr.NewInternalError(context.Background(), "filtered out")
    56  
    57  func (a *Aggregator) Close() {
    58  	// Free the StatementInfo in the Grouped
    59  	for _, item := range a.Grouped {
    60  		if stmt, ok := item.(*StatementInfo); ok {
    61  			stmt.freeNoLocked()
    62  		}
    63  	}
    64  	// clean up the Grouped map
    65  	a.Grouped = make(map[interface{}]Item)
    66  	// release resources related to the context if necessary
    67  	a.ctx = nil
    68  }
    69  
    70  func (a *Aggregator) AddItem(i Item) (Item, error) {
    71  	if !a.FilterFunc(i) {
    72  		return i, ErrFilteredOut
    73  	}
    74  
    75  	group, exists := a.Grouped[i.Key(a.WindowSize)]
    76  	if !exists {
    77  		orignal_key := i.Key(a.WindowSize)
    78  		group = a.NewItemFunc(i, a.ctx)
    79  		a.Grouped[orignal_key] = group
    80  	} else {
    81  		a.UpdateFunc(a.ctx, group, i)
    82  	}
    83  	return nil, nil
    84  }
    85  
    86  func (a *Aggregator) GetResults() []Item {
    87  	results := make([]Item, 0, len(a.Grouped))
    88  	for _, group := range a.Grouped {
    89  		results = append(results, group)
    90  	}
    91  	return results
    92  }