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
|
// 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()
}
}
func GetRequestDataStore(ctx context.Context) RequestDataStore {
if req, ok := ctx.Value(RequestDataStoreKey).(*requestDataStore); ok {
return req
}
return nil
}
type requestContext struct {
context.Context
dataStore *requestDataStore
}
func (c *requestContext) Value(key any) any {
if v := c.dataStore.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)
reqCtx := &requestContext{Context: ctx, dataStore: &requestDataStore{values: make(map[any]any)}}
return reqCtx, func() {
reqCtx.dataStore.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) context.Context {
return &requestContext{Context: parentCtx, dataStore: &requestDataStore{values: make(map[any]any)}}
}
|
