gitee.com/go-spring2/spring-base@v1.1.3/knife/knife.go (about)

     1  /*
     2   * Copyright 2012-2019 the original author or 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   *      https://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 knife provides cache on `context.Context`.
    18  package knife
    19  
    20  import (
    21  	"context"
    22  	"errors"
    23  	"fmt"
    24  	"sync"
    25  )
    26  
    27  // errUninitialized returns this error when knife is uninitialized.
    28  var errUninitialized = errors.New("knife uninitialized")
    29  
    30  var ctxKey int
    31  
    32  func cache(ctx context.Context) (*sync.Map, bool) {
    33  	m, ok := ctx.Value(&ctxKey).(*sync.Map)
    34  	return m, ok
    35  }
    36  
    37  // New returns a new context.Context with a *sync.Map bound. If a *sync.Map
    38  // is already bound, then returns the input context.Context.
    39  func New(ctx context.Context) (_ context.Context, cached bool) {
    40  	if _, ok := cache(ctx); ok {
    41  		return ctx, true
    42  	}
    43  	ctx = context.WithValue(ctx, &ctxKey, new(sync.Map))
    44  	return ctx, false
    45  }
    46  
    47  // Load returns the value stored for a key, or nil if no value is present.
    48  // If knife is uninitialized, the error of knife uninitialized will be returned.
    49  func Load(ctx context.Context, key string) (interface{}, error) {
    50  	m, ok := cache(ctx)
    51  	if !ok {
    52  		return nil, errUninitialized
    53  	}
    54  	v, _ := m.Load(key)
    55  	return v, nil
    56  }
    57  
    58  // Store stores the key and value in the context.Context, if the key is already
    59  // in the context.Context, the error of duplicate key will be returned.
    60  // If knife is uninitialized, the error of knife uninitialized will be returned.
    61  func Store(ctx context.Context, key string, val interface{}) error {
    62  	m, ok := cache(ctx)
    63  	if !ok {
    64  		return errUninitialized
    65  	}
    66  	if _, loaded := m.LoadOrStore(key, val); loaded {
    67  		return fmt.Errorf("duplicate key %q", key)
    68  	}
    69  	return nil
    70  }
    71  
    72  // LoadOrStore returns the existing value for the key if present. Otherwise, it
    73  // stores and returns the given value. The loaded result is true if the value was
    74  // loaded, false if stored.
    75  // If knife is uninitialized, the error of knife uninitialized will be returned.
    76  func LoadOrStore(ctx context.Context, key string, val interface{}) (actual interface{}, loaded bool, err error) {
    77  	m, ok := cache(ctx)
    78  	if !ok {
    79  		return nil, false, errUninitialized
    80  	}
    81  	actual, loaded = m.LoadOrStore(key, val)
    82  	return
    83  }
    84  
    85  // Delete deletes the key and value.
    86  func Delete(ctx context.Context, key string) {
    87  	if m, ok := cache(ctx); ok {
    88  		m.Delete(key)
    89  	}
    90  }
    91  
    92  // Range calls f sequentially for each key and value.
    93  func Range(ctx context.Context, f func(key, value interface{}) bool) {
    94  	if m, ok := cache(ctx); ok {
    95  		m.Range(f)
    96  	}
    97  }