github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/pkg/context/context.go (about)

     1  // Copyright 2019 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package context
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	"github.com/pingcap/tiflow/dm/pkg/log"
    21  )
    22  
    23  // Context is used to in dm to record some context field like
    24  // * go context
    25  // * logger.
    26  type Context struct {
    27  	Ctx    context.Context
    28  	Logger log.Logger
    29  }
    30  
    31  // Background return a nop context.
    32  func Background() *Context {
    33  	return &Context{
    34  		Ctx:    context.Background(),
    35  		Logger: log.L(),
    36  	}
    37  }
    38  
    39  // NewContext return a new Context.
    40  func NewContext(ctx context.Context, logger log.Logger) *Context {
    41  	return &Context{
    42  		Ctx:    ctx,
    43  		Logger: logger,
    44  	}
    45  }
    46  
    47  // WithContext set go context.
    48  func (c *Context) WithContext(ctx context.Context) *Context {
    49  	return &Context{
    50  		Ctx:    ctx,
    51  		Logger: c.Logger,
    52  	}
    53  }
    54  
    55  // WithTimeout sets a timeout associated context.
    56  func (c *Context) WithTimeout(timeout time.Duration) (*Context, context.CancelFunc) {
    57  	ctx, cancel := context.WithTimeout(c.Ctx, timeout)
    58  	return &Context{
    59  		Ctx:    ctx,
    60  		Logger: c.Logger,
    61  	}, cancel
    62  }
    63  
    64  // Context returns real context.
    65  func (c *Context) Context() context.Context {
    66  	return c.Ctx
    67  }
    68  
    69  // WithLogger set logger.
    70  func (c *Context) WithLogger(logger log.Logger) *Context {
    71  	return &Context{
    72  		Ctx:    c.Ctx,
    73  		Logger: logger,
    74  	}
    75  }
    76  
    77  // L returns real logger.
    78  func (c *Context) L() log.Logger {
    79  	return c.Logger
    80  }