sigs.k8s.io/cluster-api-provider-azure@v1.14.3/util/tele/corr_id.go (about)

     1  /*
     2  Copyright 2021 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package tele
    18  
    19  import (
    20  	"context"
    21  
    22  	"github.com/go-logr/logr"
    23  	"github.com/google/uuid"
    24  )
    25  
    26  // CorrIDKey is the type of the key used to store correlation
    27  // IDs in context.Contexts.
    28  type CorrIDKey string
    29  
    30  // CorrIDKeyVal is the key used to store the correlation ID in
    31  // context.Contexts, HTTP headers, and other similar locations.
    32  const CorrIDKeyVal CorrIDKey = "x-ms-correlation-request-id"
    33  
    34  // CorrID is a correlation ID that the cluster API provider
    35  // sends with all API requests to Azure. Do not create one
    36  // of these manually. Instead, use the CtxWithCorrelationID function
    37  // to create one of these within a context.Context.
    38  type CorrID string
    39  
    40  // ctxWithCorrID creates a CorrID and creates a new context.Context
    41  // with the new CorrID in it. It returns the _new_ context and the
    42  // newly created CorrID. If there was a problem creating the correlation
    43  // ID, the new context will not have the correlation ID in it and the
    44  // returned CorrID will be the empty string.After you call this function, prefer to
    45  // use the newly created context over the old one. Common usage is
    46  // below:
    47  //
    48  //	ctx := context.Background()
    49  //	ctx, newCorrID := ctxWithCorrID(ctx)
    50  //	fmt.Println("new corr ID: ", newCorrID)
    51  //	doSomething(ctx)
    52  func ctxWithCorrID(ctx context.Context) (context.Context, CorrID) {
    53  	if currentCorrIDIface := ctx.Value(CorrIDKeyVal); currentCorrIDIface != nil {
    54  		currentCorrID, ok := currentCorrIDIface.(CorrID)
    55  		if ok {
    56  			return ctx, currentCorrID
    57  		}
    58  	}
    59  	uid, err := uuid.NewRandom()
    60  	if err != nil {
    61  		return nil, CorrID("")
    62  	}
    63  	newCorrID := CorrID(uid.String())
    64  	ctx = context.WithValue(ctx, CorrIDKeyVal, newCorrID)
    65  	return ctx, newCorrID
    66  }
    67  
    68  // CorrIDFromCtx attempts to fetch a correlation ID from the given
    69  // context.Context. If none exists, returns an empty CorrID and false.
    70  // Otherwise returns the CorrID value and true.
    71  func CorrIDFromCtx(ctx context.Context) (CorrID, bool) {
    72  	currentCorrIDIface := ctx.Value(CorrIDKeyVal)
    73  	if currentCorrIDIface == nil {
    74  		return CorrID(""), false
    75  	}
    76  
    77  	if corrID, ok := currentCorrIDIface.(CorrID); ok {
    78  		return corrID, ok
    79  	}
    80  
    81  	return CorrID(""), false
    82  }
    83  
    84  // corrIDLogger attempts to fetch the correlation ID from the
    85  // given ctx using CorrIDFromCtx. If it finds one, this function
    86  // uses lggr.WithValues to return a new logr.Logger with the
    87  // correlation ID in it.
    88  func corrIDLogger(ctx context.Context, lggr logr.Logger) logr.Logger {
    89  	corrID, ok := CorrIDFromCtx(ctx)
    90  	if ok {
    91  		lggr = lggr.WithValues(string(CorrIDKeyVal), string(corrID))
    92  	}
    93  	return lggr
    94  }