├── README.md ├── context.go ├── context_test.go ├── go.mod ├── go.sum └── runtime └── goroutine.go /README.md: -------------------------------------------------------------------------------- 1 | # context 2 | 3 | Package context is a proof of concept implementation of **scoped context**, 4 | proposed in [this blog post](https://posener.github.io/goroutine-scoped-context). 5 | 6 | This library should not be used for production code. 7 | 8 | #### Usage 9 | 10 | The context package should be imported from `github.com/posener/context`. 11 | 12 | ```diff 13 | import ( 14 | - "context" 15 | + "github.com/posener/context" 16 | ) 17 | ``` 18 | 19 | Since this implementation does not involve changes to the runtime, 20 | the goroutine context must be initialized. 21 | 22 | ```diff 23 | func main() { 24 | + context.Init() 25 | // Go code goes here. 26 | } 27 | ``` 28 | 29 | Functions should not anymore receive the context in the first argument. 30 | They should get it from the goroutine scope. 31 | 32 | ```diff 33 | -func foo(ctx context.Context) { 34 | +func foo() { 35 | + ctx := context.Get() 36 | // Use context. 37 | } 38 | ``` 39 | 40 | Applying context to a scope: 41 | 42 | ```go 43 | unset := context.Set(ctx) 44 | // ctx is applied until unset is called, or a deeper `Set` call. 45 | unset() 46 | ``` 47 | 48 | Or: 49 | 50 | ```go 51 | defer context.Set(ctx)() 52 | // ctx is applied until the end of the function or a deeper `Set` call. 53 | ``` 54 | 55 | Invoking goroutines should be done with `context.Go` or `context.GoCtx` 56 | 57 | Running a new goroutine with the current stored context: 58 | 59 | ```diff 60 | -go foo() 61 | +context.Go(foo) 62 | ``` 63 | 64 | More complected functions: 65 | 66 | ```diff 67 | -go foo(1, "hello") 68 | +context.Go(func() { foo(1, "hello") }) 69 | ``` 70 | 71 | Running a goroutine with a new context: 72 | 73 | ```go 74 | // `ctx` is the context that we want to have in the invoked goroutine 75 | context.GoCtx(ctx, foo) 76 | ``` 77 | 78 | `context.TODO` should not be used anymore: 79 | 80 | ```diff 81 | -f(context.TODO()) 82 | +f(context.Get()) 83 | ``` 84 | 85 | ## Sub Packages 86 | 87 | * [runtime](./runtime) 88 | 89 | 90 | --- 91 | 92 | Created by [goreadme](https://github.com/apps/goreadme) 93 | -------------------------------------------------------------------------------- /context.go: -------------------------------------------------------------------------------- 1 | /* 2 | Package context is a proof of concept implementation of **scoped context**, 3 | proposed in (this blog post) https://posener.github.io/goroutine-scoped-context. 4 | 5 | This library should not be used for production code. 6 | 7 | Usage 8 | 9 | The context package should be imported from `github.com/posener/context`. 10 | 11 | import ( 12 | - "context" 13 | + "github.com/posener/context" 14 | ) 15 | 16 | Since this implementation does not involve changes to the runtime, 17 | the goroutine context must be initialized. 18 | 19 | func main() { 20 | + context.Init() 21 | // Go code goes here. 22 | } 23 | 24 | Functions should not anymore receive the context in the first argument. 25 | They should get it from the goroutine scope. 26 | 27 | -func foo(ctx context.Context) { 28 | +func foo() { 29 | + ctx := context.Get() 30 | // Use context. 31 | } 32 | 33 | Applying context to a scope: 34 | 35 | unset := context.Set(ctx) 36 | // ctx is applied until unset is called, or a deeper `Set` call. 37 | unset() 38 | 39 | Or: 40 | 41 | defer context.Set(ctx)() 42 | // ctx is applied until the end of the function or a deeper `Set` call. 43 | 44 | Invoking goroutines should be done with `context.Go` or `context.GoCtx` 45 | 46 | Running a new goroutine with the current stored context: 47 | 48 | -go foo() 49 | +context.Go(foo) 50 | 51 | More complected functions: 52 | 53 | -go foo(1, "hello") 54 | +context.Go(func() { foo(1, "hello") }) 55 | 56 | Running a goroutine with a new context: 57 | 58 | // `ctx` is the context that we want to have in the invoked goroutine 59 | context.GoCtx(ctx, foo) 60 | 61 | `context.TODO` should not be used anymore: 62 | 63 | -f(context.TODO()) 64 | +f(context.Get()) 65 | */ 66 | package context 67 | 68 | import ( 69 | stdctx "context" 70 | "sync" 71 | 72 | "github.com/posener/context/runtime" 73 | ) 74 | 75 | type ( 76 | Context = stdctx.Context 77 | CancelFunc = stdctx.CancelFunc 78 | ) 79 | 80 | var ( 81 | WithCancel = stdctx.WithCancel 82 | WithTimeout = stdctx.WithTimeout 83 | WithDeadline = stdctx.WithDeadline 84 | 85 | Background = stdctx.Background 86 | 87 | DeadlineExceeded = stdctx.DeadlineExceeded 88 | Canceled = stdctx.Canceled 89 | ) 90 | 91 | var ( 92 | // storage is used instead of goroutine local storage to 93 | // store goroutine(ID) to Context mapping. 94 | storage map[uint64][]Context 95 | // mutex for locking the storage map. 96 | mu sync.RWMutex 97 | ) 98 | 99 | func init() { 100 | storage = make(map[uint64][]Context) 101 | } 102 | 103 | // peek simulates fetching of context from goroutine local storage 104 | // It gets the context from `storage` map according to the current 105 | // goroutine ID. 106 | // If the goroutine ID is not in the map, it panic. This case 107 | // may occur when a user did not use the `context.Go` or `context.GoCtx` 108 | // to invoke a goroutine. 109 | // Note: real goroutine local storage won't need the implemented locking 110 | // exists in this implementation, since the storage won't be accessible from 111 | // different goroutines. 112 | func peek() Context { 113 | id := runtime.GID() 114 | mu.RLock() 115 | defer mu.RUnlock() 116 | stack := storage[id] 117 | if stack == nil { 118 | panic("goroutine ran without using context.Go or context.GoCtx") 119 | } 120 | return stack[len(stack)-1] 121 | } 122 | 123 | // push simulates storing of context in the goroutine local storage. 124 | // It gets the context to push to the context stack, and returns a pop function. 125 | // Note: real goroutine local storage won't need the implemented locking 126 | // exists in this implementation, since the storage won't be accessible from 127 | // different goroutines. 128 | func push(ctx Context) func() { 129 | id := runtime.GID() 130 | mu.Lock() 131 | defer mu.Unlock() 132 | storage[id] = append(storage[id], ctx) 133 | size := len(storage[id]) 134 | return func() { pop(id, size) } 135 | } 136 | 137 | // pop simulates removal of a context from the thread local storage. 138 | // If the stack is emptied, it will be removed from the storage map. 139 | // Note: real goroutine local storage won't need the implemented locking 140 | // exists in this implementation, since the storage won't be accessible from 141 | // different goroutines. 142 | func pop(id uint64, stackSize int) { 143 | mu.Lock() 144 | defer mu.Unlock() 145 | if len(storage[id]) != stackSize { 146 | if len(storage[id]) < stackSize { 147 | panic("multiple call for unset") 148 | } 149 | panic("there are contexts that should be unset before") 150 | } 151 | storage[id] = storage[id][:len(storage[id])-1] 152 | // Remove the stack from the map if it was emptied 153 | if len(storage[id]) == 0 { 154 | delete(storage, id) 155 | } 156 | } 157 | 158 | // Init creates the first background context in a program. 159 | // it should be called once, in the beginning of the main 160 | // function or in init() function. 161 | // It returns the created context. 162 | // All following goroutine invocations should be replaced 163 | // by context.Go or context.GoCtx. 164 | // 165 | // Note: 166 | // This function won't be needed in the real implementation. 167 | func Init() Context { 168 | ctx := Background() 169 | push(ctx) 170 | return ctx 171 | } 172 | 173 | // Get gets the context of the current goroutine 174 | // It may panic if the current go routine did not ran with 175 | // context.Go or context.GoCtx. 176 | // 177 | // Note: 178 | // This function won't panic in the real implementation. 179 | func Get() Context { 180 | return peek() 181 | } 182 | 183 | // Set creates a context scope. 184 | // It returns an "unset" function that should invoked at the 185 | // end of this context scope. In any case, it must be invoked, 186 | // exactly once, and in the right order. 187 | func Set(ctx Context) func() { 188 | return push(ctx) 189 | } 190 | 191 | // Go invokes f in a new goroutine and takes care of propagating 192 | // the current context to the created goroutine. 193 | // It may panic if the current goroutine was not invoked with 194 | // context.Go or context.GoCtx. 195 | // 196 | // Note: 197 | // In the real implementation, this should be the behavior 198 | // of the `go` keyword. It will also won't panic. 199 | func Go(f func()) { 200 | GoCtx(peek(), f) 201 | } 202 | 203 | // GoCtx invokes f in a new goroutine with the given context. 204 | // 205 | // Note: 206 | // In the real implementation, accepting the context argument 207 | // should be incorporated into the behavior of the `go` keyword. 208 | func GoCtx(ctx Context, f func()) { 209 | go func() { 210 | pop := push(ctx) 211 | defer pop() 212 | f() 213 | }() 214 | } 215 | -------------------------------------------------------------------------------- /context_test.go: -------------------------------------------------------------------------------- 1 | package context 2 | 3 | import ( 4 | "sync" 5 | "testing" 6 | "time" 7 | 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | type testContext int 12 | 13 | func (testContext) Deadline() (deadline time.Time, ok bool) { return } 14 | func (testContext) Done() <-chan struct{} { return make(<-chan struct{}) } 15 | func (testContext) Err() error { return nil } 16 | func (testContext) Value(key interface{}) interface{} { return nil } 17 | 18 | func TestSet(t *testing.T) { 19 | t.Parallel() 20 | ctx1 := Init() 21 | ctx2 := testContext(2) 22 | ctx3 := testContext(3) 23 | 24 | unset := Set(ctx2) 25 | 26 | var wg sync.WaitGroup 27 | wg.Add(4) 28 | 29 | Go(func() { 30 | assert.Equal(t, Get(), ctx2) 31 | wg.Done() 32 | }) 33 | 34 | GoCtx(ctx2, func() { 35 | assert.Equal(t, Get(), ctx2) 36 | wg.Done() 37 | }) 38 | 39 | GoCtx(ctx3, func() { 40 | assert.Equal(t, Get(), ctx3) 41 | wg.Done() 42 | }) 43 | 44 | unset() 45 | 46 | Go(func() { 47 | assert.Equal(t, Get(), ctx1) 48 | wg.Done() 49 | }) 50 | 51 | wg.Wait() 52 | } 53 | 54 | func TestSetNested(t *testing.T) { 55 | t.Parallel() 56 | ctx1 := Init() 57 | ctx2 := testContext(2) 58 | ctx3 := testContext(3) 59 | 60 | assert.Equal(t, Get(), ctx1) 61 | unset2 := Set(ctx2) 62 | assert.Equal(t, Get(), ctx2) 63 | unset3 := Set(ctx3) 64 | assert.Equal(t, Get(), ctx3) 65 | unset3() 66 | assert.Equal(t, Get(), ctx2) 67 | unset2() 68 | assert.Equal(t, Get(), ctx1) 69 | } 70 | 71 | func TestFunctionScope(t *testing.T) { 72 | t.Parallel() 73 | ctx1 := Init() 74 | ctx2 := testContext(2) 75 | 76 | func() { 77 | assert.Equal(t, Get(), ctx1) 78 | defer Set(ctx2)() 79 | assert.Equal(t, Get(), ctx2) 80 | }() 81 | 82 | assert.Equal(t, Get(), ctx1) 83 | } 84 | 85 | func TestPanic(t *testing.T) { 86 | t.Parallel() 87 | Init() 88 | 89 | t.Run("Using context.Get inside non-context goroutine", func(t *testing.T) { 90 | assert.Panics(t, func() { Get() }) 91 | }) 92 | 93 | t.Run("Using context.Go inside non-context goroutine", func(t *testing.T) { 94 | assert.Panics(t, func() { Go(func() {}) }) 95 | }) 96 | 97 | t.Run("Invoking unset twice", func(t *testing.T) { 98 | unset := Set(testContext(1)) 99 | unset() 100 | assert.Panics(t, unset) 101 | }) 102 | 103 | t.Run("Invoking unset unordered", func(t *testing.T) { 104 | unset1 := Set(testContext(1)) 105 | unset2 := Set(testContext(2)) 106 | assert.Panics(t, unset1) 107 | unset2() 108 | unset1() 109 | }) 110 | } 111 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/posener/context 2 | 3 | require ( 4 | github.com/KyleBanks/depth v1.2.1 // indirect 5 | github.com/davecgh/go-spew v1.1.1 // indirect 6 | github.com/pmezard/go-difflib v1.0.0 // indirect 7 | github.com/stretchr/testify v1.2.2 8 | ) 9 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= 2 | github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE= 3 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 4 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 5 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 6 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 7 | github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= 8 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= 9 | -------------------------------------------------------------------------------- /runtime/goroutine.go: -------------------------------------------------------------------------------- 1 | package runtime 2 | 3 | import ( 4 | "bytes" 5 | rt "runtime" 6 | "strconv" 7 | "sync" 8 | ) 9 | 10 | const bufSize = 64 11 | 12 | var bufPool = sync.Pool{New: func() interface{} { return make([]byte, bufSize) }} 13 | 14 | // GID returns the current goroutine ID 15 | func GID() uint64 { 16 | b := bufPool.Get().([]byte) 17 | defer bufPool.Put(b[:bufSize]) 18 | b = b[:rt.Stack(b, false)] 19 | b = bytes.TrimPrefix(b, []byte("goroutine ")) 20 | b = b[:bytes.IndexByte(b, ' ')] 21 | n, _ := strconv.ParseUint(string(b), 10, 64) 22 | return n 23 | } 24 | --------------------------------------------------------------------------------