github.com/kaiya/goutils@v1.0.1-0.20230226104005-4ae4a4dc3688/kvctx/kvctx.go (about)

     1  package kvctx
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  type kindValueCtx string
     9  
    10  const (
    11  	ctxKey kindValueCtx = "dataValueCtx"
    12  )
    13  
    14  type dataValueCtxEntry struct {
    15  	data map[string]interface{}
    16  	mu   sync.Mutex
    17  }
    18  
    19  func WithContext(ctx context.Context) context.Context {
    20  	return context.WithValue(ctx, ctxKey, &dataValueCtxEntry{
    21  		data: make(map[string]interface{}),
    22  	})
    23  }
    24  
    25  func Add(ctx context.Context, key string, value interface{}) {
    26  	v := ctx.Value(ctxKey)
    27  	if v == nil {
    28  		return
    29  	}
    30  	dataEntry, ok := v.(*dataValueCtxEntry)
    31  	if !ok {
    32  		return
    33  	}
    34  	dataEntry.mu.Lock()
    35  	defer dataEntry.mu.Unlock()
    36  	dataEntry.data[key] = value
    37  }
    38  
    39  func Get(ctx context.Context, key string) interface{} {
    40  	v := ctx.Value(ctxKey)
    41  	if v == nil {
    42  		return ctx
    43  	}
    44  	dataEntry, ok := v.(*dataValueCtxEntry)
    45  	if !ok {
    46  		return ctx
    47  	}
    48  	dataEntry.mu.Lock()
    49  	defer dataEntry.mu.Unlock()
    50  	value, ok := dataEntry.data[key]
    51  	if !ok {
    52  		return ctx
    53  	}
    54  	return value
    55  }