github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/ruler/mapper.go (about)

     1  package ruler
     2  
     3  import (
     4  	"crypto/md5"
     5  	"net/url"
     6  	"path/filepath"
     7  	"sort"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/go-kit/log/level"
    11  	"github.com/prometheus/prometheus/pkg/rulefmt"
    12  	"github.com/spf13/afero"
    13  	"gopkg.in/yaml.v3"
    14  )
    15  
    16  // mapper is designed to enusre the provided rule sets are identical
    17  // to the on-disk rules tracked by the prometheus manager
    18  type mapper struct {
    19  	Path string // Path specifies the directory in which rule files will be mapped.
    20  
    21  	FS     afero.Fs
    22  	logger log.Logger
    23  }
    24  
    25  func newMapper(path string, logger log.Logger) *mapper {
    26  	m := &mapper{
    27  		Path:   path,
    28  		FS:     afero.NewOsFs(),
    29  		logger: logger,
    30  	}
    31  	m.cleanup()
    32  
    33  	return m
    34  }
    35  
    36  func (m *mapper) cleanupUser(userID string) {
    37  	dirPath := filepath.Join(m.Path, userID)
    38  	err := m.FS.RemoveAll(dirPath)
    39  	if err != nil {
    40  		level.Warn(m.logger).Log("msg", "unable to remove user directory", "path", dirPath, "err", err)
    41  	}
    42  }
    43  
    44  // cleanup removes all of the user directories in the path of the mapper
    45  func (m *mapper) cleanup() {
    46  	level.Info(m.logger).Log("msg", "cleaning up mapped rules directory", "path", m.Path)
    47  
    48  	users, err := m.users()
    49  	if err != nil {
    50  		level.Error(m.logger).Log("msg", "unable to read rules directory", "path", m.Path, "err", err)
    51  		return
    52  	}
    53  
    54  	for _, u := range users {
    55  		m.cleanupUser(u)
    56  	}
    57  }
    58  
    59  func (m *mapper) users() ([]string, error) {
    60  	var result []string
    61  
    62  	dirs, err := afero.ReadDir(m.FS, m.Path)
    63  	for _, u := range dirs {
    64  		if u.IsDir() {
    65  			result = append(result, u.Name())
    66  		}
    67  	}
    68  
    69  	return result, err
    70  }
    71  
    72  func (m *mapper) MapRules(user string, ruleConfigs map[string][]rulefmt.RuleGroup) (bool, []string, error) {
    73  	anyUpdated := false
    74  	filenames := []string{}
    75  
    76  	// user rule files will be stored as `/<path>/<userid>/<encoded filename>`
    77  	path := filepath.Join(m.Path, user)
    78  	err := m.FS.MkdirAll(path, 0777)
    79  	if err != nil {
    80  		return false, nil, err
    81  	}
    82  
    83  	// write all rule configs to disk
    84  	for filename, groups := range ruleConfigs {
    85  		// Store the encoded file name to better handle `/` characters
    86  		encodedFileName := url.PathEscape(filename)
    87  		fullFileName := filepath.Join(path, encodedFileName)
    88  
    89  		fileUpdated, err := m.writeRuleGroupsIfNewer(groups, fullFileName)
    90  		if err != nil {
    91  			return false, nil, err
    92  		}
    93  		filenames = append(filenames, fullFileName)
    94  		anyUpdated = anyUpdated || fileUpdated
    95  	}
    96  
    97  	// and clean any up that shouldn't exist
    98  	existingFiles, err := afero.ReadDir(m.FS, path)
    99  	if err != nil {
   100  		return false, nil, err
   101  	}
   102  
   103  	for _, existingFile := range existingFiles {
   104  		fullFileName := filepath.Join(path, existingFile.Name())
   105  
   106  		// Ensure the namespace is decoded from a url path encoding to see if it is still required
   107  		decodedNamespace, err := url.PathUnescape(existingFile.Name())
   108  		if err != nil {
   109  			level.Warn(m.logger).Log("msg", "unable to remove rule file on disk", "file", fullFileName, "err", err)
   110  			continue
   111  		}
   112  
   113  		ruleGroups := ruleConfigs[string(decodedNamespace)]
   114  
   115  		if ruleGroups == nil {
   116  			err = m.FS.Remove(fullFileName)
   117  			if err != nil {
   118  				level.Warn(m.logger).Log("msg", "unable to remove rule file on disk", "file", fullFileName, "err", err)
   119  			}
   120  			anyUpdated = true
   121  		}
   122  	}
   123  
   124  	return anyUpdated, filenames, nil
   125  }
   126  
   127  func (m *mapper) writeRuleGroupsIfNewer(groups []rulefmt.RuleGroup, filename string) (bool, error) {
   128  	sort.Slice(groups, func(i, j int) bool {
   129  		return groups[i].Name > groups[j].Name
   130  	})
   131  
   132  	rgs := rulefmt.RuleGroups{Groups: groups}
   133  
   134  	d, err := yaml.Marshal(&rgs)
   135  	if err != nil {
   136  		return false, err
   137  	}
   138  
   139  	_, err = m.FS.Stat(filename)
   140  	if err == nil {
   141  		current, err := afero.ReadFile(m.FS, filename)
   142  		if err != nil {
   143  			return false, err
   144  		}
   145  		newHash := md5.New()
   146  		currentHash := md5.New()
   147  
   148  		// bailout if there is no update
   149  		if string(currentHash.Sum(current)) == string(newHash.Sum(d)) {
   150  			return false, nil
   151  		}
   152  	}
   153  
   154  	level.Info(m.logger).Log("msg", "updating rule file", "file", filename)
   155  	err = afero.WriteFile(m.FS, filename, d, 0777)
   156  	if err != nil {
   157  		return false, err
   158  	}
   159  
   160  	return true, nil
   161  }