github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/fs/disk.go (about)

     1  // Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package fsutils
     4  
     5  import (
     6  	"bytes"
     7  	"encoding/json"
     8  	"github.com/iwind/TeaGo/Tea"
     9  	"math"
    10  	"os"
    11  	"time"
    12  )
    13  
    14  const diskSpeedDataFile = "disk.speed.json"
    15  
    16  type DiskSpeedCache struct {
    17  	Speed   Speed   `json:"speed"`
    18  	SpeedMB float64 `json:"speedMB"`
    19  }
    20  
    21  // CheckDiskWritingSpeed test disk writing speed
    22  func CheckDiskWritingSpeed() (speedMB float64, err error) {
    23  	var tempDir = os.TempDir()
    24  	if len(tempDir) == 0 {
    25  		tempDir = "/tmp"
    26  	}
    27  
    28  	const filename = "edge-disk-writing-test.data"
    29  	var path = tempDir + "/" + filename
    30  	_ = os.Remove(path) // always try to delete the file
    31  
    32  	fp, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
    33  	if err != nil {
    34  		return 0, err
    35  	}
    36  
    37  	var isClosed bool
    38  	defer func() {
    39  		if !isClosed {
    40  			_ = fp.Close()
    41  		}
    42  
    43  		_ = os.Remove(path)
    44  	}()
    45  
    46  	var data = bytes.Repeat([]byte{'A'}, 16<<20)
    47  	var before = time.Now()
    48  	_, err = fp.Write(data)
    49  	if err != nil {
    50  		return 0, err
    51  	}
    52  
    53  	err = fp.Sync()
    54  	if err != nil {
    55  		return 0, err
    56  	}
    57  
    58  	err = fp.Close()
    59  	if err != nil {
    60  		return 0, err
    61  	}
    62  
    63  	var costSeconds = time.Since(before).Seconds()
    64  	speedMB = float64(len(data)) / (1 << 20) / costSeconds
    65  	speedMB = math.Ceil(speedMB/10) * 10
    66  
    67  	isClosed = true
    68  
    69  	return
    70  }
    71  
    72  // CheckDiskIsFast check disk is 'fast' disk to write
    73  func CheckDiskIsFast() (speedMB float64, isFast bool, err error) {
    74  	speedMB, err = CheckDiskWritingSpeed()
    75  	if err != nil {
    76  		return
    77  	}
    78  
    79  	isFast = speedMB > 150
    80  
    81  	if speedMB <= DiskSpeedMB {
    82  		return
    83  	}
    84  
    85  	if speedMB > 1000 {
    86  		DiskSpeed = SpeedExtremelyFast
    87  	} else if speedMB > 150 {
    88  		DiskSpeed = SpeedFast
    89  	} else if speedMB > 60 {
    90  		DiskSpeed = SpeedLow
    91  	} else {
    92  		DiskSpeed = SpeedExtremelySlow
    93  	}
    94  
    95  	DiskSpeedMB = speedMB
    96  
    97  	// write to local file
    98  	cacheData, jsonErr := json.Marshal(&DiskSpeedCache{
    99  		Speed:   DiskSpeed,
   100  		SpeedMB: DiskSpeedMB,
   101  	})
   102  	if jsonErr == nil {
   103  		_ = os.WriteFile(Tea.Root+"/data/"+diskSpeedDataFile, cacheData, 0666)
   104  	}
   105  
   106  	return
   107  }