github.com/cloudwego/kitex@v0.9.0/pkg/logid/streaming.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo 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 logid
    18  
    19  import (
    20  	"context"
    21  )
    22  
    23  // stream log id is used to identify the stream
    24  // the client should always generate a stream id, set it to its context and pass it to the server
    25  // via the header transmeta.HTTPStreamLogID; the server should get the stream id from the header,
    26  // or generate a new one if not found in the header.
    27  type keyTypeStreamLogID string
    28  
    29  const ctxKeyStreamLogID keyTypeStreamLogID = "stream-log-id"
    30  
    31  // NewCtxWithStreamLogID returns a new context with the provided log ID added as a value.
    32  //
    33  // Parameters:
    34  //   - ctx: The original context.
    35  //   - logID: The log ID to be added.
    36  //
    37  // Return:
    38  //   - context.Context: The new context with the log ID added.
    39  func NewCtxWithStreamLogID(ctx context.Context, logID string) context.Context {
    40  	return context.WithValue(ctx, ctxKeyStreamLogID, logID)
    41  }
    42  
    43  // GetStreamLogID returns the log ID from the context if it exists.
    44  //
    45  // It takes a `context.Context` as a parameter.
    46  // It returns a `string` which is the log ID, or an empty string if it doesn't exist.
    47  func GetStreamLogID(ctx context.Context) string {
    48  	if logID, ok := ctx.Value(ctxKeyStreamLogID).(string); ok {
    49  		return logID
    50  	}
    51  	return ""
    52  }
    53  
    54  // GenerateStreamLogID generates a stream log ID using the provided context.
    55  //
    56  // ctx: The context used to generate the log ID.
    57  // Returns: The generated stream log ID as a string.
    58  //
    59  // Note: ctx is for generating the log id, but the log id will not be added to the context.
    60  // You should call `NewCtxWithStreamLogID` to add the log id to the context.
    61  func GenerateStreamLogID(ctx context.Context) string {
    62  	return logIDGenerator(ctx)
    63  }