11// Copyright (c) Microsoft Corporation.
22// Licensed under the MIT License.
33
4+ using System . Collections . Concurrent ;
45using AzureMcp . Services . Interfaces ;
56using Microsoft . Extensions . Caching . Memory ;
67
@@ -9,29 +10,99 @@ namespace AzureMcp.Services.Caching;
910public class CacheService ( IMemoryCache memoryCache ) : ICacheService
1011{
1112 private readonly IMemoryCache _memoryCache = memoryCache ;
13+ private static readonly ConcurrentDictionary < string , HashSet < string > > _groupKeys = new ( ) ;
1214
13- public ValueTask < T ? > GetAsync < T > ( string key , TimeSpan ? expiration = null )
15+ public ValueTask < T ? > GetAsync < T > ( string group , string key , TimeSpan ? expiration = null )
1416 {
15- return _memoryCache . TryGetValue ( key , out T ? value ) ? new ValueTask < T ? > ( value ) : default ;
17+ string cacheKey = GetGroupKey ( group , key ) ;
18+ return _memoryCache . TryGetValue ( cacheKey , out T ? value ) ? new ValueTask < T ? > ( value ) : default ;
1619 }
1720
18- public ValueTask SetAsync < T > ( string key , T data , TimeSpan ? expiration = null )
21+ public ValueTask SetAsync < T > ( string group , string key , T data , TimeSpan ? expiration = null )
1922 {
2023 if ( data == null )
2124 return default ;
2225
26+ string cacheKey = GetGroupKey ( group , key ) ;
27+
2328 var options = new MemoryCacheEntryOptions
2429 {
2530 AbsoluteExpirationRelativeToNow = expiration
2631 } ;
2732
28- _memoryCache . Set ( key , data , options ) ;
33+ _memoryCache . Set ( cacheKey , data , options ) ;
34+
35+ // Track the key in the group
36+ _groupKeys . AddOrUpdate (
37+ group ,
38+ new HashSet < string > { key } ,
39+ ( _ , keys ) =>
40+ {
41+ keys . Add ( key ) ;
42+ return keys ;
43+ } ) ;
44+
2945 return default ;
3046 }
3147
32- public ValueTask DeleteAsync ( string key )
48+ public ValueTask DeleteAsync ( string group , string key )
3349 {
34- _memoryCache . Remove ( key ) ;
50+ string cacheKey = GetGroupKey ( group , key ) ;
51+ _memoryCache . Remove ( cacheKey ) ;
52+
53+ // Remove from group tracking
54+ if ( _groupKeys . TryGetValue ( group , out var keys ) )
55+ {
56+ keys . Remove ( key ) ;
57+ }
58+
3559 return default ;
3660 }
61+
62+ public ValueTask < IEnumerable < string > > GetGroupKeysAsync ( string group )
63+ {
64+ if ( _groupKeys . TryGetValue ( group , out var keys ) )
65+ {
66+ return new ValueTask < IEnumerable < string > > ( keys . AsEnumerable ( ) ) ;
67+ }
68+
69+ return new ValueTask < IEnumerable < string > > ( Array . Empty < string > ( ) ) ;
70+ }
71+
72+ public ValueTask ClearAsync ( )
73+ {
74+ // Clear all items from the memory cache
75+ if ( _memoryCache is MemoryCache memoryCache )
76+ {
77+ memoryCache . Compact ( 1.0 ) ;
78+ }
79+
80+ // Clear all group tracking
81+ _groupKeys . Clear ( ) ;
82+
83+ return default ;
84+ }
85+
86+ public ValueTask ClearGroupAsync ( string group )
87+ {
88+ // If this group doesn't exist, nothing to do
89+ if ( ! _groupKeys . TryGetValue ( group , out var keys ) )
90+ {
91+ return default ;
92+ }
93+
94+ // Remove each key in the group from the cache
95+ foreach ( var key in keys )
96+ {
97+ string cacheKey = GetGroupKey ( group , key ) ;
98+ _memoryCache . Remove ( cacheKey ) ;
99+ }
100+
101+ // Remove the group from tracking
102+ _groupKeys . TryRemove ( group , out _ ) ;
103+
104+ return default ;
105+ }
106+
107+ private static string GetGroupKey ( string group , string key ) => $ "{ group } :{ key } ";
37108}
0 commit comments