github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/internal/timing/context.go (about)

     1  // Copyright 2021 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package timing
     6  
     7  import (
     8  	"context"
     9  )
    10  
    11  type key int // unexported context.Context key type to avoid collisions with other packages
    12  
    13  const (
    14  	logKey          key = iota // key used for attaching a Log to a context.Context
    15  	currentStageKey            // key used for attaching a current Stage to a context.Context
    16  )
    17  
    18  // NewContext returns a new context that carries l and its root stage as
    19  // the current stage.
    20  func NewContext(ctx context.Context, l *Log) context.Context {
    21  	ctx = context.WithValue(ctx, logKey, l)
    22  	ctx = context.WithValue(ctx, currentStageKey, l.Root)
    23  	return ctx
    24  }
    25  
    26  // FromContext returns the Log and the current Stage stored in ctx, if any.
    27  func FromContext(ctx context.Context) (*Log, *Stage, bool) {
    28  	l, ok := ctx.Value(logKey).(*Log)
    29  	if !ok {
    30  		return nil, nil, false
    31  	}
    32  	s, ok := ctx.Value(currentStageKey).(*Stage)
    33  	if !ok {
    34  		return nil, nil, false
    35  	}
    36  	return l, s, true
    37  }
    38  
    39  // Start starts and returns a new Stage named name within the Log attached
    40  // to ctx. If no Log is attached to ctx, nil is returned. It is safe to call Close
    41  // on a nil stage.
    42  //
    43  // Example usage to report the time used until the end of the current function:
    44  //
    45  //	ctx, st := timing.Start(ctx, "my_stage")
    46  //	defer st.End()
    47  func Start(ctx context.Context, name string) (context.Context, *Stage) {
    48  	_, s, ok := FromContext(ctx)
    49  	if !ok {
    50  		return ctx, nil
    51  	}
    52  	c := s.StartChild(name)
    53  	return context.WithValue(ctx, currentStageKey, c), c
    54  }