context.go

/*
 * Copyright 2024 hopeio. All rights reserved.
 * Licensed under the MIT License that can be found in the LICENSE file.
 * @Created by jyb
 */

package context

import (
	"context"
	"github.com/google/uuid"
	"go.opentelemetry.io/otel/trace"
	"time"
)

type Context struct {
	ctx      context.Context
	rootSpan trace.Span
	traceID  string
}

func New(ctx context.Context) *Context {
	var traceId string
	var rootSpan trace.Span
	if ctx != nil {
		span := trace.SpanFromContext(ctx)
		if spanContext := span.SpanContext(); spanContext.IsValid() {
			traceId = spanContext.TraceID().String()
			rootSpan = span
		}
	} else {
		ctx = context.Background()
	}
	if traceId == "" || rootSpan == nil {
		ctx, rootSpan = StartSpan(ctx, "")
		if spanContext := rootSpan.SpanContext(); spanContext.IsValid() {
			traceId = spanContext.TraceID().String()
		} else {
			traceId = uuid.New().String()
		}
	}

	return &Context{
		ctx:      ctx,
		rootSpan: rootSpan,
		traceID:  traceId,
	}
}

func (c *Context) TraceID() string {
	return c.traceID
}

type ctxKey struct{}

func (c *Context) Wrapper() context.Context {
	return context.WithValue(c.ctx, ctxKey{}, c)
}

func WrapperKey() ctxKey {
	return ctxKey{}
}

func FromContextValue(ctx context.Context) *Context {
	if ctx == nil {
		return New(nil)
	}

	ctxi := ctx.Value(ctxKey{})
	c, ok := ctxi.(*Context)
	if !ok {
		c = New(ctx)
	}
	c.ctx = ctx
	return c
}

func (c *Context) Base() context.Context {
	return c.ctx
}

func (c *Context) SetBase(ctx context.Context) {
	c.ctx = ctx
}

func (c *Context) RootSpan() trace.Span {
	return c.rootSpan
}

func (c *Context) WithTimeout(timeout time.Duration) context.CancelFunc {
	var cancel context.CancelFunc
	c.ctx, cancel = context.WithTimeout(c.ctx, timeout)
	return cancel
}

func (c *Context) WithTimeoutCause(timeout time.Duration, cause error) context.CancelFunc {
	var cancel context.CancelFunc
	c.ctx, cancel = context.WithTimeoutCause(c.ctx, timeout, cause)
	return cancel
}

func (c *Context) WithCancel() context.CancelFunc {
	var cancel context.CancelFunc
	c.ctx, cancel = context.WithCancel(c.ctx)
	return cancel
}

func (c *Context) WithoutCancel() {
	c.ctx = context.WithoutCancel(c.ctx)
}

func (c *Context) WithValue(key, val any) {
	c.ctx = context.WithValue(c.ctx, key, val)
}

func (c *Context) WithCancelCause() context.CancelCauseFunc {
	var cancel context.CancelCauseFunc
	c.ctx, cancel = context.WithCancelCause(c.ctx)
	return cancel
}

func (c *Context) WithDeadline(d time.Time) context.CancelFunc {
	var cancel context.CancelFunc
	c.ctx, cancel = context.WithDeadline(c.ctx, d)
	return cancel
}

func (c *Context) WithDeadlineCause(d time.Time, cause error) context.CancelFunc {
	var cancel context.CancelFunc
	c.ctx, cancel = context.WithDeadlineCause(c.ctx, d, cause)
	return cancel
}