optimize the perf and support more features
This commit is contained in:
Lei Xue
@@ -0,0 +1,255 @@
|
||||
/*
|
||||
Copyright 2024 The GoStor Authors All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
// Package numa provides NUMA-aware utilities for multi-socket systems.
|
||||
// This package enables memory allocation optimization and thread binding
|
||||
// for better performance on NUMA architectures.
|
||||
package numa
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"runtime"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// NodeID represents a NUMA node identifier
|
||||
type NodeID int
|
||||
|
||||
// NodeInfo contains information about a NUMA node
|
||||
type NodeInfo struct {
|
||||
ID NodeID
|
||||
CPUs []int // CPU cores on this node
|
||||
TotalMemory uint64 // Total memory in bytes
|
||||
FreeMemory uint64 // Free memory in bytes
|
||||
DistanceToNode []uint32 // Distance to other nodes (lower is closer)
|
||||
}
|
||||
|
||||
// Topology represents the NUMA topology of the system
|
||||
type Topology struct {
|
||||
Nodes map[NodeID]*NodeInfo
|
||||
NumNodes int
|
||||
CPUToNodeMap map[int]NodeID
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
var (
|
||||
globalTopology *Topology
|
||||
globalTopologyOnce sync.Once
|
||||
numaAvailable bool
|
||||
)
|
||||
|
||||
// Available returns true if NUMA support is available on this system
|
||||
func Available() bool {
|
||||
return numaAvailable
|
||||
}
|
||||
|
||||
// GetTopology returns the NUMA topology of the system
|
||||
func GetTopology() *Topology {
|
||||
globalTopologyOnce.Do(func() {
|
||||
globalTopology = detectTopology()
|
||||
})
|
||||
return globalTopology
|
||||
}
|
||||
|
||||
// detectTopology detects the NUMA topology of the system
|
||||
// This is a placeholder that will be implemented per-platform
|
||||
func detectTopology() *Topology {
|
||||
topology := &Topology{
|
||||
Nodes: make(map[NodeID]*NodeInfo),
|
||||
CPUToNodeMap: make(map[int]NodeID),
|
||||
}
|
||||
|
||||
// Try to detect using platform-specific methods
|
||||
if err := detectLinuxTopology(topology); err != nil {
|
||||
// Fall back to single-node topology
|
||||
topology.NumNodes = 1
|
||||
topology.Nodes[0] = &NodeInfo{
|
||||
ID: 0,
|
||||
CPUs: makeRange(0, runtime.NumCPU()),
|
||||
TotalMemory: 0, // Unknown
|
||||
FreeMemory: 0, // Unknown
|
||||
}
|
||||
for i := 0; i < runtime.NumCPU(); i++ {
|
||||
topology.CPUToNodeMap[i] = 0
|
||||
}
|
||||
numaAvailable = false
|
||||
} else {
|
||||
numaAvailable = topology.NumNodes > 1
|
||||
}
|
||||
|
||||
return topology
|
||||
}
|
||||
|
||||
// makeRange creates a slice of integers from start to end
|
||||
func makeRange(start, end int) []int {
|
||||
result := make([]int, end-start)
|
||||
for i := range result {
|
||||
result[i] = start + i
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetNodeForCPU returns the NUMA node ID for a given CPU
|
||||
func (t *Topology) GetNodeForCPU(cpu int) (NodeID, bool) {
|
||||
t.mu.RLock()
|
||||
defer t.mu.RUnlock()
|
||||
node, ok := t.CPUToNodeMap[cpu]
|
||||
return node, ok
|
||||
}
|
||||
|
||||
// GetNode returns information about a specific NUMA node
|
||||
func (t *Topology) GetNode(id NodeID) (*NodeInfo, bool) {
|
||||
t.mu.RLock()
|
||||
defer t.mu.RUnlock()
|
||||
node, ok := t.Nodes[id]
|
||||
return node, ok
|
||||
}
|
||||
|
||||
// GetCurrentNode returns the NUMA node of the current thread
|
||||
func GetCurrentNode() (NodeID, error) {
|
||||
return getCurrentNodeImpl()
|
||||
}
|
||||
|
||||
// PreferredNode represents a preferred NUMA node for memory allocation
|
||||
type PreferredNode struct {
|
||||
nodeID NodeID
|
||||
}
|
||||
|
||||
// SetPreferredNode sets the preferred NUMA node for the current thread
|
||||
func SetPreferredNode(node NodeID) (*PreferredNode, error) {
|
||||
return setPreferredNodeImpl(node)
|
||||
}
|
||||
|
||||
// Revert restores the previous NUMA policy
|
||||
func (p *PreferredNode) Revert() error {
|
||||
return revertPreferredNodeImpl(p)
|
||||
}
|
||||
|
||||
// MemoryPolicy represents memory allocation policies
|
||||
type MemoryPolicy int
|
||||
|
||||
const (
|
||||
// MPDefault uses the default memory policy
|
||||
MPDefault MemoryPolicy = iota
|
||||
// MPBind binds memory allocation to specific nodes
|
||||
MPBind
|
||||
// MPPreferred prefers memory allocation from specific nodes
|
||||
MPPreferred
|
||||
// MPInterleave interleaves memory across nodes
|
||||
MPInterleave
|
||||
)
|
||||
|
||||
// SetMemoryPolicy sets the memory policy for the current thread
|
||||
func SetMemoryPolicy(policy MemoryPolicy, nodes []NodeID) error {
|
||||
return setMemoryPolicyImpl(policy, nodes)
|
||||
}
|
||||
|
||||
// AllocateOnNode allocates memory on a specific NUMA node
|
||||
func AllocateOnNode(size int, node NodeID) ([]byte, error) {
|
||||
return allocateOnNodeImpl(size, node)
|
||||
}
|
||||
|
||||
// LocalAlloc allocates memory on the local NUMA node
|
||||
func LocalAlloc(size int) ([]byte, error) {
|
||||
node, err := GetCurrentNode()
|
||||
if err != nil {
|
||||
// Fall back to regular allocation
|
||||
return make([]byte, size), nil
|
||||
}
|
||||
return AllocateOnNode(size, node)
|
||||
}
|
||||
|
||||
// NodeLocalPool is a memory pool that allocates from a specific NUMA node
|
||||
type NodeLocalPool struct {
|
||||
nodeID NodeID
|
||||
pool sync.Pool
|
||||
size int
|
||||
}
|
||||
|
||||
// NewNodeLocalPool creates a new NUMA-local memory pool
|
||||
func NewNodeLocalPool(size int, node NodeID) *NodeLocalPool {
|
||||
return &NodeLocalPool{
|
||||
nodeID: node,
|
||||
size: size,
|
||||
pool: sync.Pool{
|
||||
New: func() interface{} {
|
||||
buf, err := AllocateOnNode(size, node)
|
||||
if err != nil {
|
||||
// Fall back to regular allocation
|
||||
return make([]byte, size)
|
||||
}
|
||||
return buf
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a buffer from the pool
|
||||
func (p *NodeLocalPool) Get() []byte {
|
||||
return p.pool.Get().([]byte)
|
||||
}
|
||||
|
||||
// Put returns a buffer to the pool
|
||||
func (p *NodeLocalPool) Put(buf []byte) {
|
||||
if buf != nil && len(buf) >= p.size {
|
||||
p.pool.Put(buf[:p.size])
|
||||
}
|
||||
}
|
||||
|
||||
// Close releases all resources associated with the pool
|
||||
func (p *NodeLocalPool) Close() error {
|
||||
// In Go, sync.Pool doesn't have a Close method
|
||||
// The memory will be garbage collected eventually
|
||||
return nil
|
||||
}
|
||||
|
||||
// NodeScheduler schedules tasks on specific NUMA nodes
|
||||
type NodeScheduler struct {
|
||||
topology *Topology
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewNodeScheduler creates a new NUMA-aware scheduler
|
||||
func NewNodeScheduler() *NodeScheduler {
|
||||
return &NodeScheduler{
|
||||
topology: GetTopology(),
|
||||
}
|
||||
}
|
||||
|
||||
// ScheduleOnNode schedules a function to run on a specific NUMA node
|
||||
func (s *NodeScheduler) ScheduleOnNode(node NodeID, fn func()) error {
|
||||
nodeInfo, ok := s.topology.GetNode(node)
|
||||
if !ok {
|
||||
return fmt.Errorf("NUMA node %d not found", node)
|
||||
}
|
||||
|
||||
if len(nodeInfo.CPUs) == 0 {
|
||||
return fmt.Errorf("NUMA node %d has no CPUs", node)
|
||||
}
|
||||
|
||||
return scheduleOnNodeImpl(nodeInfo.CPUs[0], fn)
|
||||
}
|
||||
|
||||
// GetPreferredNodeForCurrentThread returns the preferred NUMA node
|
||||
// based on current thread's affinity
|
||||
func GetPreferredNodeForCurrentThread() NodeID {
|
||||
return getPreferredNodeForCurrentThreadImpl()
|
||||
}
|
||||
|
||||
// NumNodes returns the number of NUMA nodes in the system
|
||||
func NumNodes() int {
|
||||
return GetTopology().NumNodes
|
||||
}
|
||||
Reference in New Issue
Block a user