1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
// Copyright 2024 The Gitea Authors. All rights reserved.
// SPDX-License-Identifier: MIT
package reqctx
import (
"context"
"io"
"sync"
"code.gitea.io/gitea/modules/process"
)
type ContextDataProvider interface {
GetData() ContextData
}
type ContextData map[string]any
func (ds ContextData) GetData() ContextData {
return ds
}
func (ds ContextData) MergeFrom(other ContextData) ContextData {
for k, v := range other {
ds[k] = v
}
return ds
}
// RequestDataStore is a short-lived context-related object that is used to store request-specific data.
type RequestDataStore interface {
GetData() ContextData
SetContextValue(k, v any)
GetContextValue(key any) any
AddCleanUp(f func())
AddCloser(c io.Closer)
}
type requestDataStoreKeyType struct{}
var RequestDataStoreKey requestDataStoreKeyType
type requestDataStore struct {
data ContextData
mu sync.RWMutex
values map[any]any
cleanUpFuncs []func()
}
func (r *requestDataStore) GetContextValue(key any) any {
if key == RequestDataStoreKey {
return r
}
r.mu.RLock()
defer r.mu.RUnlock()
return r.values[key]
}
func (r *requestDataStore) SetContextValue(k, v any) {
r.mu.Lock()
r.values[k] = v
r.mu.Unlock()
}
// GetData and the underlying ContextData are not thread-safe, callers should ensure thread-safety.
func (r *requestDataStore) GetData() ContextData {
if r.data == nil {
r.data = make(ContextData)
}
return r.data
}
func (r *requestDataStore) AddCleanUp(f func()) {
r.mu.Lock()
r.cleanUpFuncs = append(r.cleanUpFuncs, f)
r.mu.Unlock()
}
func (r *requestDataStore) AddCloser(c io.Closer) {
r.AddCleanUp(func() { _ = c.Close() })
}
func (r *requestDataStore) cleanUp() {
for _, f := range r.cleanUpFuncs {
f()
}
}
type RequestContext interface {
context.Context
RequestDataStore
}
func FromContext(ctx context.Context) RequestContext {
if rc, ok := ctx.(RequestContext); ok {
return rc
}
// here we must use the current ctx and the underlying store
// the current ctx guarantees that the ctx deadline/cancellation/values are respected
// the underlying store guarantees that the request-specific data is available
if store := GetRequestDataStore(ctx); store != nil {
return &requestContext{Context: ctx, RequestDataStore: store}
}
return nil
}
func GetRequestDataStore(ctx context.Context) RequestDataStore {
if req, ok := ctx.Value(RequestDataStoreKey).(*requestDataStore); ok {
return req
}
return nil
}
type requestContext struct {
context.Context
RequestDataStore
}
func (c *requestContext) Value(key any) any {
if v := c.GetContextValue(key); v != nil {
return v
}
return c.Context.Value(key)
}
func NewRequestContext(parentCtx context.Context, profDesc string) (_ context.Context, finished func()) {
ctx, _, processFinished := process.GetManager().AddTypedContext(parentCtx, profDesc, process.RequestProcessType, true)
store := &requestDataStore{values: make(map[any]any)}
reqCtx := &requestContext{Context: ctx, RequestDataStore: store}
return reqCtx, func() {
store.cleanUp()
processFinished()
}
}
// NewRequestContextForTest creates a new RequestContext for testing purposes
// It doesn't add the context to the process manager, nor do cleanup
func NewRequestContextForTest(parentCtx context.Context) RequestContext {
return &requestContext{Context: parentCtx, RequestDataStore: &requestDataStore{values: make(map[any]any)}}
}
|
