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

     1  // Copyright 2023 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 util
    15  
    16  import (
    17  	"context"
    18  	"math"
    19  	"time"
    20  
    21  	"github.com/KimMachineGun/automemlimit/memlimit"
    22  	"github.com/pingcap/log"
    23  	"github.com/pingcap/tidb/pkg/util/memory"
    24  	"github.com/pingcap/tiflow/pkg/errors"
    25  	"github.com/shirou/gopsutil/v3/mem"
    26  	"go.uber.org/zap"
    27  )
    28  
    29  const memoryMax uint64 = math.MaxUint64
    30  
    31  // GetMemoryLimit gets the memory limit of current process based on cgroup.
    32  // If the cgourp is not set or memory.max is set to max, returns the available
    33  // memory of host.
    34  func GetMemoryLimit() (uint64, error) {
    35  	totalMemory, err := memlimit.FromCgroup()
    36  	if err != nil || totalMemory == memoryMax {
    37  		log.Info("no cgroup memory limit", zap.Error(err))
    38  		totalMemory, err = memory.MemTotal()
    39  		if err != nil {
    40  			return 0, errors.Trace(err)
    41  		}
    42  	}
    43  	return totalMemory, nil
    44  }
    45  
    46  // CheckMemoryUsage checks if the memory usage is less than the limit.
    47  func CheckMemoryUsage(limit float64) (bool, error) {
    48  	stat, err := mem.VirtualMemory()
    49  	if err != nil {
    50  		return false, err
    51  	}
    52  
    53  	log.Info("check memory usage", zap.Any("memory", stat))
    54  	return stat.UsedPercent < limit, nil
    55  }
    56  
    57  // WaitMemoryAvailable waits until the memory usage is less than the limit.
    58  func WaitMemoryAvailable(ctx context.Context, limit float64, timeout time.Duration) error {
    59  	ticker := time.NewTicker(time.Second * 5)
    60  	timeoutTimer := time.NewTimer(timeout)
    61  	for {
    62  		select {
    63  		case <-ctx.Done():
    64  			return errors.WrapError(errors.ErrWaitFreeMemoryTimeout, ctx.Err())
    65  		case <-ticker.C:
    66  			hasFreeMemory, err := CheckMemoryUsage(limit)
    67  			if err != nil {
    68  				return err
    69  			}
    70  			if hasFreeMemory {
    71  				return nil
    72  			}
    73  		case <-timeoutTimer.C:
    74  			return errors.ErrWaitFreeMemoryTimeout.GenWithStackByArgs()
    75  		}
    76  	}
    77  }