diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1b8e008 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019 Zhandos Zhylkaidar + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..44d1e7d --- /dev/null +++ b/README.md @@ -0,0 +1,11 @@ +# LGORE +CPU Affinity for Go goroutines (Unix systems only) + +## Warnings +Use at your own risk + +## Usage +TODO + +## Why +TODO \ No newline at end of file diff --git a/environment.go b/environment.go new file mode 100644 index 0000000..ff98c9a --- /dev/null +++ b/environment.go @@ -0,0 +1,153 @@ +package gofine + +import ( + "errors" + "sync" + + "golang.org/x/sys/unix" +) + +const maxNumCPUs = 1 << 10 + +var invalidLgoreId = errors.New("Invalid lgore id") + +// Config environment config +type Config struct { + // Specifies if we should pre-occupy all available cores + // + // TODO ignored for now + OccupyAll bool + + // Specifies which cores should not be used as lgores + // + // There should be at least one core present in this slice. + // Each value should be an index in range [0, NumCPUs), + // where NumCPUs can be found from runtime.NumCPU(). + ReserveCores []int +} + +// Environment manages all the lgores +type Environment struct { + mu sync.Mutex + original unix.CPUSet + available unix.CPUSet + lgores []*lgore +} + +// InitDefault initializes `env` with default configuration +// +// Default config sets `OccupyAll` to `false` and adds core 0 for reserve +func (env *Environment) InitDefault() error { + defaultConf := Config{} + defaultConf.OccupyAll = false + defaultConf.ReserveCores = append(defaultConf.ReserveCores, 0) + + return env.Init(defaultConf) +} + +// Init initializes environment and lgores +func (env *Environment) Init(conf Config) error { + if len(conf.ReserveCores) == 0 { + return errors.New("Should reserve at least one lgore") + } + + // save original cpu affinity + err := unix.SchedGetaffinity(0, &env.original) + if err != nil { + return err + } + + if env.original.Count() <= 1 { + return errors.New("Not enough logical cores, should be greater than one") + } + + env.available = env.original + + // reserve cores for Go runtime + for _, coreIndex := range conf.ReserveCores { + coreId := getCoreIdByIndex(env.original, coreIndex) + if coreId < 0 { + return errors.New("Invalid reservation lgore id") + } + + env.available.Clear(coreId) + } + if env.available.Count() == 0 { + return errors.New("No lgores left after reservation") + } + + env.initLgores() + // TODO occupy lgores if OccupyAll is true + return nil +} + +// LgoreCount returns number of available lgores +func (env *Environment) LgoreCount() int { + return env.available.Count() +} + +// GetLgoreState returns `LgoreState` of a lgore +func (env *Environment) GetLgoreState(lgoreId int) (LgoreState, error) { + if lgoreId >= len(env.lgores) { + return Invalid, invalidLgoreId + } + + return env.lgores[lgoreId].state, nil +} + +// Occupy locks calling goroutine to an lgore +// +// Goroutine is locked to OS thread, and OS thread is locked to lgore's core. +func (env *Environment) Occupy(lgoreId int) error { + if lgoreId >= len(env.lgores) { + return invalidLgoreId + } + + env.mu.Lock() + defer env.mu.Unlock() + + lg := env.lgores[lgoreId] + return lg.occupy() +} + +// Release releases the lgore +// +// This function should be called from the same goroutine that called `Occupy`. +// Lgore becomes available, and the locked OS thread allowed to run on any core again. +func (env *Environment) Release(lgoreId int) error { + if lgoreId >= len(env.lgores) { + return invalidLgoreId + } + + env.mu.Lock() + defer env.mu.Unlock() + + lg := env.lgores[lgoreId] + return lg.release(env.original) +} + +func (env *Environment) initLgores() { + env.lgores = make([]*lgore, env.available.Count()) + + for lgoreId := 0; lgoreId < len(env.lgores); lgoreId++ { + coreId := getCoreIdByIndex(env.available, lgoreId) + env.lgores[lgoreId] = &lgore{coreId: coreId, state: Available} + } +} + +// returns -1 if not found +func getCoreIdByIndex(cpuset unix.CPUSet, coreIndex int) int { + count := 0 + + for coreId := 0; coreId < maxNumCPUs; coreId++ { + if cpuset.IsSet(coreId) { + if coreIndex == count { + return coreId + } + + count++ + } + } + + return -1 +} diff --git a/environment_test.go b/environment_test.go new file mode 100644 index 0000000..8037ac0 --- /dev/null +++ b/environment_test.go @@ -0,0 +1,51 @@ +package gofine_test + +import ( + "testing" + + "github.com/jandos/gofine" +) + +var ( + env = &gofine.Environment{} +) + +func TestProcess(t *testing.T) { + // TODO write parallel tests + err := env.InitDefault() + if err != nil { + t.Fatal(err) + } + + lgoreCount := env.LgoreCount() + if lgoreCount <= 0 { + t.Fatal("lgore count should be greater than zero") + } + + lgoreId := 0 + err = env.Occupy(lgoreId) + if err != nil { + t.Fatal(err) + } + + state, err := env.GetLgoreState(lgoreId) + if err != nil { + t.Fatal(err) + } + if state != gofine.Busy { + t.Fatal("lgore state should be Busy") + } + + err = env.Release(lgoreId) + if err != nil { + t.Fatal(err) + } + + state, err = env.GetLgoreState(lgoreId) + if err != nil { + t.Fatal(err) + } + if state != gofine.Available { + t.Fatal("lgore state should be Available") + } +} diff --git a/examples/busyworker/main.go b/examples/busyworker/main.go new file mode 100644 index 0000000..d701e25 --- /dev/null +++ b/examples/busyworker/main.go @@ -0,0 +1,39 @@ +package main + +import ( + "log" + "sync" + + "github.com/jandos/gofine" +) + +func main() { + var env gofine.Environment + err := env.InitDefault() + if err != nil { + panic(err) + } + + log.Printf("Available worker count: %v\n", env.LgoreCount()) + + var wg sync.WaitGroup + wg.Add(1) + go func(lgoreId int) { + defer wg.Done() + err := env.Occupy(lgoreId) + if err != nil { + panic(err) + } + defer env.Release(lgoreId) + + incrementMeHard := 0 + for { + // do non-interruptible super important work + // open up htop and verify that goroutine doesn't jump around + // and runs on the specified core index + incrementMeHard++ + } + }(0) + + wg.Wait() +} diff --git a/lgore.go b/lgore.go new file mode 100644 index 0000000..1eb1ef3 --- /dev/null +++ b/lgore.go @@ -0,0 +1,53 @@ +package gofine + +import ( + "errors" + "runtime" + + "golang.org/x/sys/unix" +) + +// LgoreState represents lgore's current state +type LgoreState uint8 + +const ( + // Invalid state for non-existing lgores + Invalid LgoreState = iota + + // Available represents an lgore which can be occupied + Available + + // Busy represents an lgore which is occupied + Busy +) + +type lgore struct { + coreId int + state LgoreState +} + +func (lg *lgore) occupy() error { + if lg.state == Busy { + return errors.New("lgore is busy") + } + runtime.LockOSThread() + + var cpuset unix.CPUSet + cpuset.Set(lg.coreId) + + err := unix.SchedSetaffinity(0, &cpuset) + if err == nil { + lg.state = Busy + } + return err +} + +func (lg *lgore) release(original unix.CPUSet) error { + if lg.state == Available { + return nil + } + defer runtime.UnlockOSThread() + + lg.state = Available + return unix.SchedSetaffinity(0, &original) +} diff --git a/lgore_test.go b/lgore_test.go new file mode 100644 index 0000000..a46657b --- /dev/null +++ b/lgore_test.go @@ -0,0 +1,77 @@ +package gofine + +import ( + "os" + "runtime" + "testing" + + "golang.org/x/sys/unix" +) + +var ( + busyLgore = lgore{coreId: 0, state: Busy} + availableLgore = lgore{coreId: 0, state: Available} + original unix.CPUSet +) + +func TestOccupyBusy(t *testing.T) { + err := busyLgore.occupy() + + if err == nil { + t.Fatal("err should not be nil") + } +} + +func TestReleaseAvailable(t *testing.T) { + err := availableLgore.release(original) + + if err != nil { + t.Fatal("err should be nil") + } +} + +func TestOccupyRelease(t *testing.T) { + lgore := availableLgore + // do occupy + err := lgore.occupy() + if err != nil { + t.Fatalf("err: %v", err.Error()) + } + + if lgore.state != Busy { + t.Fatal("state should be Busy") + } + + var cpuset unix.CPUSet + unix.SchedGetaffinity(0, &cpuset) + if !cpuset.IsSet(lgore.coreId) { + t.Fatal("coreId should be set") + } + + // do release + err = lgore.release(original) + if err != nil { + t.Fatalf("err: %v", err.Error()) + } + + if lgore.state != Available { + t.Fatal("state should be Available") + } + + unix.SchedGetaffinity(0, &cpuset) + if cpuset != original { + t.Fatal("cpuset should be equal to original") + } +} + +func setup() { + runtime.LockOSThread() + defer runtime.UnlockOSThread() + unix.SchedGetaffinity(0, &original) +} + +func TestMain(m *testing.M) { + setup() + code := m.Run() + os.Exit(code) +}