go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/tsmon/context.go (about)

     1  // Copyright 2016 The LUCI 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 tsmon
    16  
    17  import (
    18  	"context"
    19  
    20  	"go.chromium.org/luci/common/tsmon/monitor"
    21  	"go.chromium.org/luci/common/tsmon/store"
    22  	"go.chromium.org/luci/common/tsmon/target"
    23  )
    24  
    25  // WithState returns a new context holding the given State instance.
    26  func WithState(ctx context.Context, s *State) context.Context {
    27  	return context.WithValue(ctx, stateKey, s)
    28  }
    29  
    30  // WithFakes returns a new context holding a new State with a fake store and a
    31  // fake monitor.
    32  func WithFakes(ctx context.Context) (context.Context, *store.Fake, *monitor.Fake) {
    33  	s := &store.Fake{}
    34  	m := &monitor.Fake{}
    35  	return WithState(ctx, &State{
    36  		store:                        s,
    37  		monitor:                      m,
    38  		invokeGlobalCallbacksOnFlush: 1,
    39  	}), s, m
    40  }
    41  
    42  // WithDummyInMemory returns a new context holding a new State with a new in-
    43  // memory store and a fake monitor.
    44  func WithDummyInMemory(ctx context.Context) (context.Context, *monitor.Fake) {
    45  	m := &monitor.Fake{}
    46  	return WithState(ctx, &State{
    47  		store:                        store.NewInMemory(&target.Task{}),
    48  		monitor:                      m,
    49  		invokeGlobalCallbacksOnFlush: 1,
    50  	}), m
    51  }
    52  
    53  // stateFromContext returns a State instance from the given context, if there
    54  // is one; otherwise the default State is returned.
    55  func stateFromContext(ctx context.Context) *State {
    56  	if ret := ctx.Value(stateKey); ret != nil {
    57  		return ret.(*State)
    58  	}
    59  
    60  	return globalState
    61  }
    62  
    63  type key string
    64  
    65  var stateKey = key("tsmon.State")