@@ -6,6 +6,7 @@ package cache
66import (
77 "context"
88 "sync"
9+ "time"
910
1011 "code.gitea.io/gitea/modules/log"
1112)
@@ -14,65 +15,151 @@ import (
1415// This is useful for caching data that is expensive to calculate and is likely to be
1516// used multiple times in a request.
1617type cacheContext struct {
17- ctx context.Context
18- data map [any ]map [any ]any
19- lock sync.RWMutex
18+ data map [any ]map [any ]any
19+ lock sync.RWMutex
20+ created time.Time
21+ discard bool
2022}
2123
2224func (cc * cacheContext ) Get (tp , key any ) any {
2325 cc .lock .RLock ()
2426 defer cc .lock .RUnlock ()
25- if cc .data [tp ] == nil {
26- return nil
27- }
2827 return cc.data [tp ][key ]
2928}
3029
3130func (cc * cacheContext ) Put (tp , key , value any ) {
3231 cc .lock .Lock ()
3332 defer cc .lock .Unlock ()
34- if cc .data [tp ] == nil {
35- cc .data [tp ] = make (map [any ]any )
33+
34+ if cc .discard {
35+ return
36+ }
37+
38+ d := cc .data [tp ]
39+ if d == nil {
40+ d = make (map [any ]any )
41+ cc .data [tp ] = d
3642 }
37- cc. data [ tp ] [key ] = value
43+ d [key ] = value
3844}
3945
4046func (cc * cacheContext ) Delete (tp , key any ) {
4147 cc .lock .Lock ()
4248 defer cc .lock .Unlock ()
43- if cc .data [tp ] == nil {
44- return
45- }
4649 delete (cc .data [tp ], key )
4750}
4851
52+ func (cc * cacheContext ) Discard () {
53+ cc .lock .Lock ()
54+ defer cc .lock .Unlock ()
55+ cc .data = nil
56+ cc .discard = true
57+ }
58+
59+ func (cc * cacheContext ) isDiscard () bool {
60+ cc .lock .RLock ()
61+ defer cc .lock .RUnlock ()
62+ return cc .discard
63+ }
64+
65+ // cacheContextLifetime is the max lifetime of cacheContext.
66+ // Since cacheContext is used to cache data in a request level context, 10s is enough.
67+ // If a cacheContext is used more than 10s, it's probably misuse.
68+ const cacheContextLifetime = 10 * time .Second
69+
70+ var timeNow = time .Now
71+
72+ func (cc * cacheContext ) Expired () bool {
73+ return timeNow ().Sub (cc .created ) > cacheContextLifetime
74+ }
75+
4976var cacheContextKey = struct {}{}
5077
78+ /*
79+ Since there are both WithCacheContext and WithNoCacheContext,
80+ it may be confusing when there is nesting.
81+
82+ Some cases to explain the design:
83+
84+ When:
85+ - A, B or C means a cache context.
86+ - A', B' or C' means a discard cache context.
87+ - ctx means context.Backgrand().
88+ - A(ctx) means a cache context with ctx as the parent context.
89+ - B(A(ctx)) means a cache context with A(ctx) as the parent context.
90+ - With is alias of WithCacheContext.
91+ - WithNo is alias of WithNoCacheContext.
92+
93+ So:
94+ - With(ctx) -> A(ctx)
95+ - With(With(ctx)) -> A(ctx), not B(A(ctx)), always reuse parent cache context if possible.
96+ - With(With(With(ctx))) -> A(ctx), not C(B(A(ctx))), ditto.
97+ - WithNo(ctx) -> ctx, not A'(ctx), don't create new cache context if we don't have to.
98+ - WithNo(With(ctx)) -> A'(ctx)
99+ - WithNo(WithNo(With(ctx))) -> A'(ctx), not B'(A'(ctx)), don't create new cache context if we don't have to.
100+ - With(WithNo(With(ctx))) -> B(A'(ctx)), not A(ctx), never reuse a discard cache context.
101+ - WithNo(With(WithNo(With(ctx)))) -> B'(A'(ctx))
102+ - With(WithNo(With(WithNo(With(ctx))))) -> C(B'(A'(ctx))), so there's always only one not-discard cache context.
103+ */
104+
51105func WithCacheContext (ctx context.Context ) context.Context {
106+ if c , ok := ctx .Value (cacheContextKey ).(* cacheContext ); ok {
107+ if ! c .isDiscard () {
108+ // reuse parent context
109+ return ctx
110+ }
111+ }
52112 return context .WithValue (ctx , cacheContextKey , & cacheContext {
53- ctx : ctx ,
54- data : make ( map [ any ] map [ any ] any ),
113+ data : make ( map [ any ] map [ any ] any ) ,
114+ created : timeNow ( ),
55115 })
56116}
57117
118+ func WithNoCacheContext (ctx context.Context ) context.Context {
119+ if c , ok := ctx .Value (cacheContextKey ).(* cacheContext ); ok {
120+ // The caller want to run long-life tasks, but the parent context is a cache context.
121+ // So we should disable and clean the cache data, or it will be kept in memory for a long time.
122+ c .Discard ()
123+ return ctx
124+ }
125+
126+ return ctx
127+ }
128+
58129func GetContextData (ctx context.Context , tp , key any ) any {
59130 if c , ok := ctx .Value (cacheContextKey ).(* cacheContext ); ok {
131+ if c .Expired () {
132+ // The warning means that the cache context is misused for long-life task,
133+ // it can be resolved with WithNoCacheContext(ctx).
134+ log .Warn ("cache context is expired, may be misused for long-life tasks: %v" , c )
135+ return nil
136+ }
60137 return c .Get (tp , key )
61138 }
62- log .Warn ("cannot get cache context when getting data: %v" , ctx )
63139 return nil
64140}
65141
66142func SetContextData (ctx context.Context , tp , key , value any ) {
67143 if c , ok := ctx .Value (cacheContextKey ).(* cacheContext ); ok {
144+ if c .Expired () {
145+ // The warning means that the cache context is misused for long-life task,
146+ // it can be resolved with WithNoCacheContext(ctx).
147+ log .Warn ("cache context is expired, may be misused for long-life tasks: %v" , c )
148+ return
149+ }
68150 c .Put (tp , key , value )
69151 return
70152 }
71- log .Warn ("cannot get cache context when setting data: %v" , ctx )
72153}
73154
74155func RemoveContextData (ctx context.Context , tp , key any ) {
75156 if c , ok := ctx .Value (cacheContextKey ).(* cacheContext ); ok {
157+ if c .Expired () {
158+ // The warning means that the cache context is misused for long-life task,
159+ // it can be resolved with WithNoCacheContext(ctx).
160+ log .Warn ("cache context is expired, may be misused for long-life tasks: %v" , c )
161+ return
162+ }
76163 c .Delete (tp , key )
77164 }
78165}
0 commit comments