memory.go

package sql

import (
	"os"
	"runtime"
	"strconv"
	"sync"

	errors "gopkg.in/src-d/go-errors.v1"
)

// Disposable objects can erase all their content when they're no longer in use.
// They should not be used again after they've been disposed.
type Disposable interface {
	// Dispose the contents.
	Dispose()
}

// Freeable objects can free their memory.
type Freeable interface {
	// Free the memory.
	Free()
}

// KeyValueCache is a cache of key value pairs.
type KeyValueCache interface {
	// Put a new value in the cache.
	Put(uint64, interface{}) error
	// Get the value with the given key.
	Get(uint64) (interface{}, error)
}

// RowsCache is a cache of rows.
type RowsCache interface {
	// Add a new row to the cache. If there is no memory available, it will try to
	// free some memory. If after that there is still no memory available, it
	// will return an error and erase all the content of the cache.
	Add(Row) error
	// Get all rows.
	Get() []Row
}

// ErrNoMemoryAvailable is returned when there is no more available memory.
var ErrNoMemoryAvailable = errors.NewKind("no memory available")

const maxMemoryKey = "MAX_MEMORY"

const (
	b   = 1
	kib = 1024 * b
	mib = 1024 * kib
)

var maxMemory = func() uint64 {
	val := os.Getenv(maxMemoryKey)
	var v uint64
	if val != "" {
		var err error
		v, err = strconv.ParseUint(val, 10, 64)
		if err != nil {
			panic("MAX_MEMORY environment variable must be a number, but got: " + val)
		}
	}

	return v * uint64(mib)
}()

// Reporter is a component that gives information about the memory usage.
type Reporter interface {
	// MaxMemory returns the maximum number of memory allowed in bytes.
	MaxMemory() uint64
	// UsedMemory returns the memory in use in bytes.
	UsedMemory() uint64
}

// ProcessMemory is a reporter for the memory used by the process and the
// maximum amount of memory allowed controlled by the MAX_MEMORY environment
// variable.
var ProcessMemory Reporter = new(processReporter)

type processReporter struct{}

func (processReporter) UsedMemory() uint64 {
	var s runtime.MemStats
	runtime.ReadMemStats(&s)
	return s.HeapInuse + s.StackInuse
}

func (processReporter) MaxMemory() uint64 { return maxMemory }

// HasAvailableMemory reports whether more memory is available to the program if
// it hasn't reached the max memory limit.
func HasAvailableMemory(r Reporter) bool {
	maxMemory := r.MaxMemory()
	if maxMemory == 0 {
		return true
	}

	return r.UsedMemory() < maxMemory
}

// MemoryManager is in charge of keeping track and managing all the components that operate
// in memory. There should only be one instance of a memory manager running at the
// same time in each process.
type MemoryManager struct {
	mu       sync.RWMutex
	reporter Reporter
	caches   map[uint64]Disposable
	token    uint64
}

// NewMemoryManager creates a new manager with the given memory reporter. If nil is given,
// then the Process reporter will be used by default.
func NewMemoryManager(r Reporter) *MemoryManager {
	if r == nil {
		r = ProcessMemory
	}

	return &MemoryManager{
		reporter: r,
		caches:   make(map[uint64]Disposable),
	}
}

// HasAvailable reports whether the memory manager has any available memory.
func (m *MemoryManager) HasAvailable() bool {
	return HasAvailableMemory(m.reporter)
}

// DisposeFunc is a function to completely erase a cache and remove it from the manager.
type DisposeFunc func()

// // NewLRUCache returns an empty LRU cache and a function to dispose it when it's
// // no longer needed.
// func (m *MemoryManager) NewLRUCache(size uint) (KeyValueCache, DisposeFunc) {
// 	c := newLRUCache(m, m.reporter, size)
// 	pos := m.addCache(c)
// 	return c, func() {
// 		c.Dispose()
// 		m.removeCache(pos)
// 	}
// }

// // NewHistoryCache returns an empty history cache and a function to dispose it when it's
// // no longer needed.
// func (m *MemoryManager) NewHistoryCache() (KeyValueCache, DisposeFunc) {
// 	c := newHistoryCache(m, m.reporter)
// 	pos := m.addCache(c)
// 	return c, func() {
// 		c.Dispose()
// 		m.removeCache(pos)
// 	}
// }

// // NewRowsCache returns an empty rows cache and a function to dispose it when it's
// // no longer needed.
// func (m *MemoryManager) NewRowsCache() (RowsCache, DisposeFunc) {
// 	c := newRowsCache(m, m.reporter)
// 	pos := m.addCache(c)
// 	return c, func() {
// 		c.Dispose()
// 		m.removeCache(pos)
// 	}
// }

func (m *MemoryManager) addCache(c Disposable) (pos uint64) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.token++
	m.caches[m.token] = c
	return m.token
}

func (m *MemoryManager) removeCache(pos uint64) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.caches, pos)

	if len(m.caches) == 0 {
		m.token = 0
	}
}

// Free the memory of all freeable caches.
func (m *MemoryManager) Free() {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, c := range m.caches {
		if f, ok := c.(Freeable); ok {
			f.Free()
		}
	}
}