github.com/matrixorigin/matrixone@v1.2.0/pkg/perfcounter/context.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 perfcounter
    16  
    17  import "context"
    18  
    19  type CounterSets = map[*CounterSet]struct{}
    20  
    21  type ctxKeyCounters struct{}
    22  
    23  var CtxKeyCounters = ctxKeyCounters{}
    24  
    25  func WithCounterSet(ctx context.Context, sets ...*CounterSet) context.Context {
    26  	// check existed
    27  	v := ctx.Value(CtxKeyCounters)
    28  	if v == nil {
    29  		v := make(CounterSets)
    30  		for _, s := range sets {
    31  			if s == nil {
    32  				panic("nil counter set")
    33  			}
    34  			v[s] = struct{}{}
    35  		}
    36  		return context.WithValue(ctx, CtxKeyCounters, v)
    37  	}
    38  
    39  	counters := v.(CounterSets)
    40  
    41  	allExist := true
    42  	for _, s := range sets {
    43  		if _, ok := counters[s]; !ok {
    44  			allExist = false
    45  			break
    46  		}
    47  	}
    48  
    49  	// if all exist already, try not to nest context too depth
    50  	if allExist {
    51  		return ctx
    52  	}
    53  
    54  	newCounters := make(CounterSets, len(counters)+1)
    55  	for counter := range counters {
    56  		newCounters[counter] = struct{}{}
    57  	}
    58  
    59  	for _, s := range sets {
    60  		if s == nil {
    61  			panic("nil counter set")
    62  		}
    63  		newCounters[s] = struct{}{}
    64  	}
    65  	return context.WithValue(ctx, CtxKeyCounters, newCounters)
    66  }
    67  
    68  func WithCounterSetFrom(ctx context.Context, fromCtx context.Context) context.Context {
    69  	if v := fromCtx.Value(CtxKeyCounters); v != nil {
    70  		var sets []*CounterSet
    71  		for set := range v.(CounterSets) {
    72  			sets = append(sets, set)
    73  		}
    74  		return WithCounterSet(ctx, sets...)
    75  	}
    76  	return ctx
    77  }