Skip to content

Commit d375458

Browse files
manusamanusa
committed
feat(kubernetes): reusable Kubernetes clients
Improve cache performance
1 parent 40ff50e commit d375458

File tree

10 files changed

+131
-159
lines changed

10 files changed

+131
-159
lines changed

‎pkg/kubernetes-mcp-server/cmd/root.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,11 @@ Kubernetes Model Context Protocol (MCP) server
2828
fmt.Println(version.Version)
2929
return
3030
}
31-
if err := mcp.NewSever().ServeStdio(); err != nil && !errors.Is(err, context.Canceled) {
31+
mcpServer, err := mcp.NewSever()
32+
if err != nil {
33+
panic(err)
34+
}
35+
if err := mcpServer.ServeStdio(); err != nil && !errors.Is(err, context.Canceled) {
3236
panic(err)
3337
}
3438
},

‎pkg/kubernetes/kubernetes.go‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@ package kubernetes
22

33
import (
44
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
5+
"k8s.io/client-go/discovery"
6+
"k8s.io/client-go/discovery/cached/memory"
7+
"k8s.io/client-go/dynamic"
8+
"k8s.io/client-go/kubernetes"
59
"k8s.io/client-go/rest"
610
"k8s.io/client-go/restmapper"
711
"k8s.io/client-go/tools/clientcmd"
@@ -11,15 +15,36 @@ import (
1115

1216
type Kubernetes struct {
1317
cfg *rest.Config
18+
clientSet *kubernetes.Clientset
19+
discoveryClient *discovery.DiscoveryClient
1420
deferredDiscoveryRESTMapper *restmapper.DeferredDiscoveryRESTMapper
21+
dynamicClient *dynamic.DynamicClient
1522
}
1623

1724
func NewKubernetes() (*Kubernetes, error) {
1825
cfg, err := resolveClientConfig()
1926
if err != nil {
2027
return nil, err
2128
}
22-
return &Kubernetes{cfg: cfg}, nil
29+
clientSet, err := kubernetes.NewForConfig(cfg)
30+
if err != nil {
31+
return nil, err
32+
}
33+
discoveryClient, err := discovery.NewDiscoveryClientForConfig(cfg)
34+
if err != nil {
35+
return nil, err
36+
}
37+
dynamicClient, err := dynamic.NewForConfig(cfg)
38+
if err != nil {
39+
return nil, err
40+
}
41+
return &Kubernetes{
42+
cfg: cfg,
43+
clientSet: clientSet,
44+
discoveryClient: discoveryClient,
45+
deferredDiscoveryRESTMapper: restmapper.NewDeferredDiscoveryRESTMapper(memory.NewMemCacheClient(discoveryClient)),
46+
dynamicClient: dynamicClient,
47+
}, nil
2348
}
2449

2550
func marshal(v any) (string, error) {
@@ -28,8 +53,14 @@ func marshal(v any) (string, error) {
2853
for i := range t {
2954
t[i].SetManagedFields(nil)
3055
}
56+
case []*unstructured.Unstructured:
57+
for i := range t {
58+
t[i].SetManagedFields(nil)
59+
}
3160
case unstructured.Unstructured:
3261
t.SetManagedFields(nil)
62+
case *unstructured.Unstructured:
63+
t.SetManagedFields(nil)
3364
}
3465
ret, err := yaml.Marshal(v)
3566
if err != nil {

‎pkg/kubernetes/pods.go‎

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@ import (
1111
"k8s.io/apimachinery/pkg/runtime/schema"
1212
"k8s.io/apimachinery/pkg/util/intstr"
1313
"k8s.io/apimachinery/pkg/util/rand"
14-
"k8s.io/client-go/dynamic"
15-
"k8s.io/client-go/kubernetes"
1614
)
1715

1816
func (k *Kubernetes) PodsListInAllNamespaces(ctx context.Context) (string, error) {
@@ -34,13 +32,8 @@ func (k *Kubernetes) PodsGet(ctx context.Context, namespace, name string) (strin
3432
}
3533

3634
func (k *Kubernetes) PodsDelete(ctx context.Context, namespace, name string) (string, error) {
37-
cs, err := kubernetes.NewForConfig(k.cfg)
38-
if err != nil {
39-
return "", err
40-
}
41-
4235
namespace = namespaceOrDefault(namespace)
43-
pod, err := cs.CoreV1().Pods(namespace).Get(ctx, name, metav1.GetOptions{})
36+
pod, err := k.clientSet.CoreV1().Pods(namespace).Get(ctx, name, metav1.GetOptions{})
4437
if err != nil {
4538
return "", err
4639
}
@@ -53,22 +46,18 @@ func (k *Kubernetes) PodsDelete(ctx context.Context, namespace, name string) (st
5346

5447
// Delete managed service
5548
if isManaged {
56-
if sl, _ := cs.CoreV1().Services(namespace).List(ctx, metav1.ListOptions{
49+
if sl, _ := k.clientSet.CoreV1().Services(namespace).List(ctx, metav1.ListOptions{
5750
LabelSelector: managedLabelSelector.String(),
5851
}); sl != nil {
5952
for _, svc := range sl.Items {
60-
_ = cs.CoreV1().Services(namespace).Delete(ctx, svc.Name, metav1.DeleteOptions{})
53+
_ = k.clientSet.CoreV1().Services(namespace).Delete(ctx, svc.Name, metav1.DeleteOptions{})
6154
}
6255
}
6356
}
6457

6558
// Delete managed Route
6659
if isManaged && k.supportsGroupVersion("route.openshift.io/v1") {
67-
dynamicClient, dErr := dynamic.NewForConfig(k.cfg)
68-
if dErr != nil {
69-
return "", dErr
70-
}
71-
routeResources := dynamicClient.
60+
routeResources := k.dynamicClient.
7261
Resource(schema.GroupVersionResource{Group: "route.openshift.io", Version: "v1", Resource: "routes"}).
7362
Namespace(namespace)
7463
if rl, _ := routeResources.List(ctx, metav1.ListOptions{
@@ -80,16 +69,13 @@ func (k *Kubernetes) PodsDelete(ctx context.Context, namespace, name string) (st
8069
}
8170

8271
}
83-
return "Pod deleted successfully", cs.CoreV1().Pods(namespace).Delete(ctx, name, metav1.DeleteOptions{})
72+
return "Pod deleted successfully",
73+
k.clientSet.CoreV1().Pods(namespace).Delete(ctx, name, metav1.DeleteOptions{})
8474
}
8575

8676
func (k *Kubernetes) PodsLog(ctx context.Context, namespace, name string) (string, error) {
87-
cs, err := kubernetes.NewForConfig(k.cfg)
88-
if err != nil {
89-
return "", err
90-
}
9177
tailLines := int64(256)
92-
req := cs.CoreV1().Pods(namespaceOrDefault(namespace)).GetLogs(name, &v1.PodLogOptions{
78+
req := k.clientSet.CoreV1().Pods(namespaceOrDefault(namespace)).GetLogs(name, &v1.PodLogOptions{
9379
TailLines: &tailLines,
9480
})
9581
res := req.Do(ctx)

‎pkg/kubernetes/resources.go‎

Lines changed: 11 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,6 @@ import (
77
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
88
"k8s.io/apimachinery/pkg/runtime/schema"
99
"k8s.io/apimachinery/pkg/util/yaml"
10-
"k8s.io/client-go/discovery"
11-
memory "k8s.io/client-go/discovery/cached"
12-
"k8s.io/client-go/dynamic"
13-
"k8s.io/client-go/restmapper"
1410
"regexp"
1511
"strings"
1612
)
@@ -23,26 +19,18 @@ const (
2319
)
2420

2521
func (k *Kubernetes) ResourcesList(ctx context.Context, gvk *schema.GroupVersionKind, namespace string) (string, error) {
26-
client, err := dynamic.NewForConfig(k.cfg)
27-
if err != nil {
28-
return "", err
29-
}
3022
gvr, err := k.resourceFor(gvk)
3123
if err != nil {
3224
return "", err
3325
}
34-
rl, err := client.Resource(*gvr).Namespace(namespace).List(ctx, metav1.ListOptions{})
26+
rl, err := k.dynamicClient.Resource(*gvr).Namespace(namespace).List(ctx, metav1.ListOptions{})
3527
if err != nil {
3628
return "", err
3729
}
3830
return marshal(rl.Items)
3931
}
4032

4133
func (k *Kubernetes) ResourcesGet(ctx context.Context, gvk *schema.GroupVersionKind, namespace, name string) (string, error) {
42-
client, err := dynamic.NewForConfig(k.cfg)
43-
if err != nil {
44-
return "", err
45-
}
4634
gvr, err := k.resourceFor(gvk)
4735
if err != nil {
4836
return "", err
@@ -51,7 +39,7 @@ func (k *Kubernetes) ResourcesGet(ctx context.Context, gvk *schema.GroupVersionK
5139
if namespaced, nsErr := k.isNamespaced(gvk); nsErr == nil && namespaced {
5240
namespace = namespaceOrDefault(namespace)
5341
}
54-
rg, err := client.Resource(*gvr).Namespace(namespace).Get(ctx, name, metav1.GetOptions{})
42+
rg, err := k.dynamicClient.Resource(*gvr).Namespace(namespace).Get(ctx, name, metav1.GetOptions{})
5543
if err != nil {
5644
return "", err
5745
}
@@ -73,10 +61,6 @@ func (k *Kubernetes) ResourcesCreateOrUpdate(ctx context.Context, resource strin
7361
}
7462

7563
func (k *Kubernetes) ResourcesDelete(ctx context.Context, gvk *schema.GroupVersionKind, namespace, name string) error {
76-
client, err := dynamic.NewForConfig(k.cfg)
77-
if err != nil {
78-
return err
79-
}
8064
gvr, err := k.resourceFor(gvk)
8165
if err != nil {
8266
return err
@@ -85,14 +69,10 @@ func (k *Kubernetes) ResourcesDelete(ctx context.Context, gvk *schema.GroupVersi
8569
if namespaced, nsErr := k.isNamespaced(gvk); nsErr == nil && namespaced {
8670
namespace = namespaceOrDefault(namespace)
8771
}
88-
return client.Resource(*gvr).Namespace(namespace).Delete(ctx, name, metav1.DeleteOptions{})
72+
return k.dynamicClient.Resource(*gvr).Namespace(namespace).Delete(ctx, name, metav1.DeleteOptions{})
8973
}
9074

9175
func (k *Kubernetes) resourcesCreateOrUpdate(ctx context.Context, resources []*unstructured.Unstructured) (string, error) {
92-
client, err := dynamic.NewForConfig(k.cfg)
93-
if err != nil {
94-
return "", err
95-
}
9676
for i, obj := range resources {
9777
gvk := obj.GroupVersionKind()
9878
gvr, rErr := k.resourceFor(&gvk)
@@ -104,24 +84,21 @@ func (k *Kubernetes) resourcesCreateOrUpdate(ctx context.Context, resources []*u
10484
if namespaced, nsErr := k.isNamespaced(&gvk); nsErr == nil && namespaced {
10585
namespace = namespaceOrDefault(namespace)
10686
}
107-
resources[i], rErr = client.Resource(*gvr).Namespace(namespace).Apply(ctx, obj.GetName(), obj, metav1.ApplyOptions{
87+
resources[i], rErr = k.dynamicClient.Resource(*gvr).Namespace(namespace).Apply(ctx, obj.GetName(), obj, metav1.ApplyOptions{
10888
FieldManager: version.BinaryName,
10989
})
11090
if rErr != nil {
11191
return "", rErr
11292
}
93+
// Clear the cache to ensure the next operation is performed on the latest exposed APIs
94+
if gvk.Kind == "CustomResourceDefinition" {
95+
k.deferredDiscoveryRESTMapper.Reset()
96+
}
11397
}
11498
return marshal(resources)
11599
}
116100

117101
func (k *Kubernetes) resourceFor(gvk *schema.GroupVersionKind) (*schema.GroupVersionResource, error) {
118-
if k.deferredDiscoveryRESTMapper == nil {
119-
d, err := discovery.NewDiscoveryClientForConfig(k.cfg)
120-
if err != nil {
121-
return nil, err
122-
}
123-
k.deferredDiscoveryRESTMapper = restmapper.NewDeferredDiscoveryRESTMapper(memory.NewMemCacheClient(d))
124-
}
125102
m, err := k.deferredDiscoveryRESTMapper.RESTMapping(schema.GroupKind{Group: gvk.Group, Kind: gvk.Kind}, gvk.Version)
126103
if err != nil {
127104
return nil, err
@@ -130,11 +107,7 @@ func (k *Kubernetes) resourceFor(gvk *schema.GroupVersionKind) (*schema.GroupVer
130107
}
131108

132109
func (k *Kubernetes) isNamespaced(gvk *schema.GroupVersionKind) (bool, error) {
133-
d, err := discovery.NewDiscoveryClientForConfig(k.cfg)
134-
if err != nil {
135-
return false, err
136-
}
137-
apiResourceList, err := d.ServerResourcesForGroupVersion(gvk.GroupVersion().String())
110+
apiResourceList, err := k.discoveryClient.ServerResourcesForGroupVersion(gvk.GroupVersion().String())
138111
if err != nil {
139112
return false, err
140113
}
@@ -147,13 +120,8 @@ func (k *Kubernetes) isNamespaced(gvk *schema.GroupVersionKind) (bool, error) {
147120
}
148121

149122
func (k *Kubernetes) supportsGroupVersion(groupVersion string) bool {
150-
d, err := discovery.NewDiscoveryClientForConfig(k.cfg)
151-
if err != nil {
123+
if _, err := k.discoveryClient.ServerResourcesForGroupVersion(groupVersion); err != nil {
152124
return false
153125
}
154-
_, err = d.ServerResourcesForGroupVersion(groupVersion)
155-
if err == nil {
156-
return true
157-
}
158-
return false
126+
return true
159127
}

‎pkg/mcp/common_test.go‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -77,23 +77,29 @@ func TestMain(m *testing.M) {
7777
}
7878

7979
type mcpContext struct {
80-
ctx context.Context
81-
tempDir string
82-
testServer *httptest.Server
83-
cancel context.CancelFunc
84-
mcpClient *client.SSEMCPClient
80+
ctx context.Context
81+
tempDir string
82+
cancel context.CancelFunc
83+
mcpServer *Server
84+
mcpHttpServer *httptest.Server
85+
mcpClient *client.SSEMCPClient
8586
}
8687

8788
func (c *mcpContext) beforeEach(t *testing.T) {
8889
var err error
8990
c.ctx, c.cancel = context.WithCancel(context.Background())
9091
c.tempDir = t.TempDir()
91-
c.withKubeConfig(nil)
92-
c.testServer = server.NewTestServer(NewSever().server)
93-
if c.mcpClient, err = client.NewSSEMCPClient(c.testServer.URL + "/sse"); err != nil {
92+
_ = os.Unsetenv("KUBECONFIG")
93+
if c.mcpServer, err = NewSever(); err != nil {
94+
t.Fatal(err)
95+
return
96+
}
97+
c.mcpHttpServer = server.NewTestServer(c.mcpServer.server)
98+
if c.mcpClient, err = client.NewSSEMCPClient(c.mcpHttpServer.URL + "/sse"); err != nil {
9499
t.Fatal(err)
95100
return
96101
}
102+
c.withKubeConfig(nil)
97103
if err = c.mcpClient.Start(c.ctx); err != nil {
98104
t.Fatal(err)
99105
return
@@ -111,7 +117,7 @@ func (c *mcpContext) beforeEach(t *testing.T) {
111117
func (c *mcpContext) afterEach() {
112118
c.cancel()
113119
_ = c.mcpClient.Close()
114-
c.testServer.Close()
120+
c.mcpHttpServer.Close()
115121
}
116122

117123
func testCase(t *testing.T, test func(c *mcpContext)) {
@@ -140,6 +146,9 @@ func (c *mcpContext) withKubeConfig(rc *rest.Config) *api.Config {
140146
kubeConfig := filepath.Join(c.tempDir, "config")
141147
_ = clientcmd.WriteToFile(*fakeConfig, kubeConfig)
142148
_ = os.Setenv("KUBECONFIG", kubeConfig)
149+
if err := c.mcpServer.reloadKubernetesClient(); err != nil {
150+
panic(err)
151+
}
143152
return fakeConfig
144153
}
145154

@@ -168,11 +177,9 @@ func (c *mcpContext) inOpenShift() func() {
168177
}`)
169178
}
170179

171-
// newKubernetesClient creates a new Kubernetes client with the current kubeconfig
180+
// newKubernetesClient creates a new Kubernetes client with the envTest kubeconfig
172181
func (c *mcpContext) newKubernetesClient() *kubernetes.Clientset {
173-
c.withEnvTest()
174-
cfg, _ := clientcmd.BuildConfigFromFlags("", clientcmd.NewDefaultPathOptions().GetDefaultFilename())
175-
return kubernetes.NewForConfigOrDie(cfg)
182+
return kubernetes.NewForConfigOrDie(envTestRestConfig)
176183
}
177184

178185
// newApiExtensionsClient creates a new ApiExtensions client with the envTest kubeconfig

‎pkg/mcp/configuration.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
"github.com/mark3labs/mcp-go/mcp"
88
)
99

10-
func (s *Sever) initConfiguration() {
10+
func (s *Server) initConfiguration() {
1111
s.server.AddTool(mcp.NewTool(
1212
"configuration_view",
1313
mcp.WithDescription("Get the current Kubernetes configuration content as a kubeconfig YAML"),

0 commit comments

Comments
 (0)