github.com/pkg/sftp@v1.13.6/allocator.go (about)

     1  package sftp
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type allocator struct {
     8  	sync.Mutex
     9  	available [][]byte
    10  	// map key is the request order
    11  	used map[uint32][][]byte
    12  }
    13  
    14  func newAllocator() *allocator {
    15  	return &allocator{
    16  		// micro optimization: initialize available pages with an initial capacity
    17  		available: make([][]byte, 0, SftpServerWorkerCount*2),
    18  		used:      make(map[uint32][][]byte),
    19  	}
    20  }
    21  
    22  // GetPage returns a previously allocated and unused []byte or create a new one.
    23  // The slice have a fixed size = maxMsgLength, this value is suitable for both
    24  // receiving new packets and reading the files to serve
    25  func (a *allocator) GetPage(requestOrderID uint32) []byte {
    26  	a.Lock()
    27  	defer a.Unlock()
    28  
    29  	var result []byte
    30  
    31  	// get an available page and remove it from the available ones.
    32  	if len(a.available) > 0 {
    33  		truncLength := len(a.available) - 1
    34  		result = a.available[truncLength]
    35  
    36  		a.available[truncLength] = nil          // clear out the internal pointer
    37  		a.available = a.available[:truncLength] // truncate the slice
    38  	}
    39  
    40  	// no preallocated slice found, just allocate a new one
    41  	if result == nil {
    42  		result = make([]byte, maxMsgLength)
    43  	}
    44  
    45  	// put result in used pages
    46  	a.used[requestOrderID] = append(a.used[requestOrderID], result)
    47  
    48  	return result
    49  }
    50  
    51  // ReleasePages marks unused all pages in use for the given requestID
    52  func (a *allocator) ReleasePages(requestOrderID uint32) {
    53  	a.Lock()
    54  	defer a.Unlock()
    55  
    56  	if used := a.used[requestOrderID]; len(used) > 0 {
    57  		a.available = append(a.available, used...)
    58  	}
    59  	delete(a.used, requestOrderID)
    60  }
    61  
    62  // Free removes all the used and available pages.
    63  // Call this method when the allocator is not needed anymore
    64  func (a *allocator) Free() {
    65  	a.Lock()
    66  	defer a.Unlock()
    67  
    68  	a.available = nil
    69  	a.used = make(map[uint32][][]byte)
    70  }
    71  
    72  func (a *allocator) countUsedPages() int {
    73  	a.Lock()
    74  	defer a.Unlock()
    75  
    76  	num := 0
    77  	for _, p := range a.used {
    78  		num += len(p)
    79  	}
    80  	return num
    81  }
    82  
    83  func (a *allocator) countAvailablePages() int {
    84  	a.Lock()
    85  	defer a.Unlock()
    86  
    87  	return len(a.available)
    88  }
    89  
    90  func (a *allocator) isRequestOrderIDUsed(requestOrderID uint32) bool {
    91  	a.Lock()
    92  	defer a.Unlock()
    93  
    94  	_, ok := a.used[requestOrderID]
    95  	return ok
    96  }