github.com/TeaOSLab/EdgeNode@v1.3.8/internal/caches/partial_ranges_queue.go (about)

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package caches
     4  
     5  import (
     6  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     7  	"github.com/TeaOSLab/EdgeNode/internal/goman"
     8  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
     9  	"github.com/TeaOSLab/EdgeNode/internal/utils/fnv"
    10  	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
    11  	memutils "github.com/TeaOSLab/EdgeNode/internal/utils/mem"
    12  	"sync"
    13  )
    14  
    15  var SharedPartialRangesQueue = NewPartialRangesQueue()
    16  
    17  func init() {
    18  	if !teaconst.IsMain {
    19  		return
    20  	}
    21  
    22  	SharedPartialRangesQueue.Start()
    23  }
    24  
    25  const partialRangesQueueSharding = 8
    26  
    27  // PartialRangesQueue ranges file writing queue
    28  type PartialRangesQueue struct {
    29  	m [partialRangesQueueSharding]map[string][]byte // { filename => data, ... }
    30  
    31  	c  chan string // filename1, ...
    32  	mu [partialRangesQueueSharding]*sync.RWMutex
    33  }
    34  
    35  // NewPartialRangesQueue Create new queue
    36  func NewPartialRangesQueue() *PartialRangesQueue {
    37  	var queueSize = 512
    38  	var memGB = memutils.SystemMemoryGB()
    39  	if memGB > 16 {
    40  		queueSize = 8 << 10
    41  	} else if memGB > 8 {
    42  		queueSize = 4 << 10
    43  	} else if memGB > 4 {
    44  		queueSize = 2 << 10
    45  	} else if memGB > 2 {
    46  		queueSize = 1 << 10
    47  	}
    48  
    49  	var m = [partialRangesQueueSharding]map[string][]byte{}
    50  	var muList = [partialRangesQueueSharding]*sync.RWMutex{}
    51  	for i := 0; i < partialRangesQueueSharding; i++ {
    52  		muList[i] = &sync.RWMutex{}
    53  		m[i] = map[string][]byte{}
    54  	}
    55  
    56  	return &PartialRangesQueue{
    57  		mu: muList,
    58  		m:  m,
    59  		c:  make(chan string, queueSize),
    60  	}
    61  }
    62  
    63  // Start the queue
    64  func (this *PartialRangesQueue) Start() {
    65  	goman.New(func() {
    66  		this.Dump()
    67  	})
    68  }
    69  
    70  // Put ranges data to filename
    71  func (this *PartialRangesQueue) Put(filename string, data []byte) {
    72  	var index = this.indexForKey(filename)
    73  
    74  	this.mu[index].Lock()
    75  	this.m[index][filename] = data
    76  	this.mu[index].Unlock()
    77  
    78  	// always wait to finish
    79  	this.c <- filename
    80  }
    81  
    82  // Get ranges data from filename
    83  func (this *PartialRangesQueue) Get(filename string) ([]byte, error) {
    84  	var index = this.indexForKey(filename)
    85  
    86  	this.mu[index].RLock()
    87  	data, ok := this.m[index][filename]
    88  	this.mu[index].RUnlock()
    89  
    90  	if ok {
    91  		return data, nil
    92  	}
    93  
    94  	return fsutils.ReadFile(filename)
    95  }
    96  
    97  // Delete ranges filename
    98  func (this *PartialRangesQueue) Delete(filename string) {
    99  	var index = this.indexForKey(filename)
   100  
   101  	this.mu[index].Lock()
   102  	delete(this.m[index], filename)
   103  	this.mu[index].Unlock()
   104  }
   105  
   106  // Dump ranges to filename from memory
   107  func (this *PartialRangesQueue) Dump() {
   108  	for filename := range this.c {
   109  		var index = this.indexForKey(filename)
   110  
   111  		this.mu[index].Lock()
   112  		data, ok := this.m[index][filename]
   113  		if ok {
   114  			delete(this.m[index], filename)
   115  		}
   116  		this.mu[index].Unlock()
   117  
   118  		if !ok || len(data) == 0 {
   119  			continue
   120  		}
   121  
   122  		err := fsutils.WriteFile(filename, data, 0666)
   123  		if err != nil {
   124  			remotelogs.Println("PARTIAL_RANGES_QUEUE", "write file '"+filename+"' failed: "+err.Error())
   125  		}
   126  	}
   127  }
   128  
   129  // Len count all files
   130  func (this *PartialRangesQueue) Len() int {
   131  	var count int
   132  
   133  	for i := 0; i < partialRangesQueueSharding; i++ {
   134  		this.mu[i].RLock()
   135  		count += len(this.m[i])
   136  		this.mu[i].RUnlock()
   137  	}
   138  
   139  	return count
   140  }
   141  
   142  func (this *PartialRangesQueue) indexForKey(filename string) int {
   143  	return int(fnv.HashString(filename) % partialRangesQueueSharding)
   144  }