github.com/go-spring/spring-base@v1.1.3/clock/now.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 clock
    18  
    19  import (
    20  	"context"
    21  	"time"
    22  
    23  	"github.com/go-spring/spring-base/knife"
    24  	"github.com/go-spring/spring-base/run"
    25  )
    26  
    27  const nowKey = "::now::"
    28  
    29  type FakeTime interface {
    30  	Now() time.Time
    31  }
    32  
    33  type fixedTime struct {
    34  	fixed time.Time
    35  }
    36  
    37  func (t *fixedTime) Now() time.Time {
    38  	return t.fixed
    39  }
    40  
    41  type baseTime struct {
    42  	base time.Time
    43  	from time.Time
    44  }
    45  
    46  func (t *baseTime) Now() time.Time {
    47  	return t.base.Add(time.Since(t.from))
    48  }
    49  
    50  // ResetTime resets the time to normal.
    51  func ResetTime(ctx context.Context) {
    52  	knife.Delete(ctx, nowKey)
    53  }
    54  
    55  // SetBaseTime sets the base time.
    56  func SetBaseTime(ctx context.Context, t time.Time) error {
    57  	return setTime(ctx, &baseTime{base: t, from: time.Now()})
    58  }
    59  
    60  // SetFixedTime sets a fixed time.
    61  func SetFixedTime(ctx context.Context, t time.Time) error {
    62  	return setTime(ctx, &fixedTime{fixed: t})
    63  }
    64  
    65  func setTime(ctx context.Context, t FakeTime) error {
    66  	ResetTime(ctx)
    67  	return knife.Store(ctx, nowKey, t)
    68  }
    69  
    70  // Now returns the current local time.
    71  func Now(ctx context.Context) time.Time {
    72  	if ctx == nil || run.NormalMode() || run.RecordMode() {
    73  		return time.Now()
    74  	}
    75  	v, err := knife.Load(ctx, nowKey)
    76  	if err != nil || v == nil {
    77  		return time.Now()
    78  	}
    79  	t, ok := v.(FakeTime)
    80  	if !ok {
    81  		return time.Now()
    82  	}
    83  	return t.Now()
    84  }