@@ -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