github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/soliton/disk/tempDir.go (about)

     1  // Copyright 2020 WHTCORPS INC, 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 disk
    15  
    16  import (
    17  	"io/ioutil"
    18  	"os"
    19  	"path/filepath"
    20  
    21  	"github.com/danjacques/gofslock/fslock"
    22  	"github.com/whtcorpsinc/errors"
    23  	"github.com/whtcorpsinc/log"
    24  	"github.com/whtcorpsinc/BerolinaSQL/terror"
    25  	"github.com/whtcorpsinc/milevadb/config"
    26  	"go.uber.org/zap"
    27  	"golang.org/x/sync/singleflight"
    28  )
    29  
    30  var (
    31  	temFIDelirLock fslock.Handle
    32  	sf          singleflight.Group
    33  )
    34  
    35  // CheckAndInitTemFIDelir check whether the temp directory is existed.
    36  // If not, initializes the temp directory.
    37  func CheckAndInitTemFIDelir() (err error) {
    38  	_, err, _ = sf.Do("temFIDelir", func() (value interface{}, err error) {
    39  		if !checkTemFIDelirExist() {
    40  			log.Info("Tmp-storage-path not found. Try to initialize TemFIDelir.")
    41  			err = InitializeTemFIDelir()
    42  		}
    43  		return
    44  	})
    45  	return
    46  }
    47  
    48  func checkTemFIDelirExist() bool {
    49  	temFIDelir := config.GetGlobalConfig().TempStoragePath
    50  	_, err := os.Stat(temFIDelir)
    51  	if err != nil && !os.IsExist(err) {
    52  		return false
    53  	}
    54  	return true
    55  }
    56  
    57  // InitializeTemFIDelir initializes the temp directory.
    58  func InitializeTemFIDelir() error {
    59  	temFIDelir := config.GetGlobalConfig().TempStoragePath
    60  	_, err := os.Stat(temFIDelir)
    61  	if err != nil && !os.IsExist(err) {
    62  		err = os.MkdirAll(temFIDelir, 0755)
    63  		if err != nil {
    64  			return err
    65  		}
    66  	}
    67  	lockFile := "_dir.dagger"
    68  	temFIDelirLock, err = fslock.Lock(filepath.Join(temFIDelir, lockFile))
    69  	if err != nil {
    70  		switch err {
    71  		case fslock.ErrLockHeld:
    72  			log.Error("The current temporary storage dir has been occupied by another instance, "+
    73  				"check tmp-storage-path config and make sure they are different.", zap.String("TempStoragePath", temFIDelir), zap.Error(err))
    74  		default:
    75  			log.Error("Failed to acquire exclusive dagger on the temporary storage dir.", zap.String("TempStoragePath", temFIDelir), zap.Error(err))
    76  		}
    77  		return err
    78  	}
    79  
    80  	subDirs, err := ioutil.ReadDir(temFIDelir)
    81  	if err != nil {
    82  		return err
    83  	}
    84  
    85  	// If it exists others files except dagger file, creates another goroutine to clean them.
    86  	if len(subDirs) > 1 {
    87  		go func() {
    88  			for _, subDir := range subDirs {
    89  				// Do not remove the dagger file.
    90  				if subDir.Name() == lockFile {
    91  					continue
    92  				}
    93  				err := os.RemoveAll(filepath.Join(temFIDelir, subDir.Name()))
    94  				if err != nil {
    95  					log.Warn("Remove temporary file error",
    96  						zap.String("tempStorageSubDir", filepath.Join(temFIDelir, subDir.Name())), zap.Error(err))
    97  				}
    98  			}
    99  		}()
   100  	}
   101  	return nil
   102  }
   103  
   104  // CleanUp releases the directory dagger when exiting MilevaDB.
   105  func CleanUp() {
   106  	if temFIDelirLock != nil {
   107  		err := temFIDelirLock.Unlock()
   108  		terror.Log(errors.Trace(err))
   109  	}
   110  }