Skip to content

Commit 85e158f

Browse files
jentingjenting
authored
Merge pull request #50 from hsiaoairplane/refactor/lister-context-logging
Use informer lister, context-aware shutdown, and standard logging
2 parents 7094b3e + 6dec9f6 commit 85e158f

2 files changed

Lines changed: 97 additions & 74 deletions

File tree

‎main.go‎

Lines changed: 68 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ import (
44
"context"
55
"flag"
66
"fmt"
7+
"log"
78
"os"
89
"os/exec"
910
"os/signal"
1011
"path/filepath"
12+
"sort"
1113
"strings"
1214
"sync"
1315
"syscall"
@@ -16,8 +18,10 @@ import (
1618
v1 "k8s.io/api/core/v1"
1719
apierrors "k8s.io/apimachinery/pkg/api/errors"
1820
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
21+
"k8s.io/apimachinery/pkg/labels"
1922
"k8s.io/client-go/informers"
2023
"k8s.io/client-go/kubernetes"
24+
corelisters "k8s.io/client-go/listers/core/v1"
2125
"k8s.io/client-go/rest"
2226
"k8s.io/client-go/tools/cache"
2327
"k8s.io/client-go/tools/clientcmd"
@@ -30,8 +34,9 @@ var (
3034
watchKey string
3135
mainContainerName string
3236

33-
clientset kubernetes.Interface
34-
lock sync.Mutex
37+
clientset kubernetes.Interface
38+
namespaceLister corelisters.NamespaceLister
39+
lock sync.Mutex
3540

3641
// restartMainContainer is the action used to reload the ConfigMap.
3742
// It is a variable so tests can override it.
@@ -53,51 +58,44 @@ func main() {
5358

5459
config, err := getKubeConfig()
5560
if err != nil {
56-
fmt.Printf("Error getting Kubernetes config: %v\n", err)
57-
os.Exit(1)
61+
log.Fatalf("Error getting Kubernetes config: %v", err)
5862
}
5963

6064
clientset, err = kubernetes.NewForConfig(config)
6165
if err != nil {
62-
fmt.Printf("Error creating Kubernetes client: %v\n", err)
63-
os.Exit(1)
66+
log.Fatalf("Error creating Kubernetes client: %v", err)
6467
}
6568

66-
// Ensure the ConfigMap is created/updated at startup
67-
updateConfigMap()
69+
// Cancel the context on SIGTERM/SIGINT for graceful shutdown.
70+
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
71+
defer stop()
6872

69-
// Set up signal handling for graceful shutdown
70-
stopCh := make(chan struct{})
71-
signalCh := make(chan os.Signal, 1)
72-
73-
// Listen for termination signals
74-
signal.Notify(signalCh, syscall.SIGTERM, syscall.SIGINT)
75-
76-
go func() {
77-
<-signalCh
78-
fmt.Println("Received termination signal. Shutting down gracefully...")
79-
close(stopCh) // Stop the namespace watcher
80-
os.Exit(0) // Exit the program
81-
}()
73+
// Start the namespace watcher and wait for its cache to sync.
74+
if err := watchNamespaces(ctx); err != nil {
75+
log.Fatalf("Error starting namespace watcher: %v", err)
76+
}
8277

83-
// Start namespace watcher
84-
go watchNamespaces(stopCh)
78+
// Ensure the ConfigMap reflects the current state once the cache is synced.
79+
// This also covers the case where no matching namespace exists yet, for
80+
// which no informer Add event would fire.
81+
updateConfigMap(ctx)
8582

86-
// Keep the program running
87-
<-stopCh
83+
// Block until a termination signal cancels the context.
84+
<-ctx.Done()
85+
log.Println("Received termination signal. Shutting down gracefully...")
8886
}
8987

9088
// getKubeConfig tries in-cluster config first, then falls back to local kubeconfig
9189
func getKubeConfig() (*rest.Config, error) {
9290
// Try in-cluster config
9391
config, err := rest.InClusterConfig()
9492
if err == nil {
95-
fmt.Println("Using in-cluster configuration")
93+
log.Println("Using in-cluster configuration")
9694
return config, nil
9795
}
9896

9997
// Fall back to local kubeconfig
100-
fmt.Println("In-cluster config failed, falling back to local kubeconfig")
98+
log.Println("In-cluster config failed, falling back to local kubeconfig")
10199
kubeconfig := filepath.Join(os.Getenv("HOME"), ".kube", "config")
102100
if envKubeconfig := os.Getenv("KUBECONFIG"); envKubeconfig != "" {
103101
kubeconfig = envKubeconfig
@@ -108,48 +106,54 @@ func getKubeConfig() (*rest.Config, error) {
108106
return nil, fmt.Errorf("failed to load kubeconfig: %w", err)
109107
}
110108

111-
fmt.Println("Using local kubeconfig")
109+
log.Println("Using local kubeconfig")
112110
return config, nil
113111
}
114112

115-
// watchNamespaces monitors namespace creation and deletion
116-
func watchNamespaces(stopCh chan struct{}) {
113+
// watchNamespaces starts an informer that monitors namespace changes and keeps
114+
// the ConfigMap in sync. It returns once the informer cache has synced; the
115+
// informer keeps running in the background until ctx is cancelled.
116+
func watchNamespaces(ctx context.Context) error {
117117
informerFactory := informers.NewSharedInformerFactoryWithOptions(clientset, time.Minute, informers.WithTweakListOptions(func(opts *metav1.ListOptions) {
118118
opts.LabelSelector = labelSelector
119119
}))
120-
namespaceInformer := informerFactory.Core().V1().Namespaces().Informer()
120+
namespaceInformer := informerFactory.Core().V1().Namespaces()
121+
namespaceLister = namespaceInformer.Lister()
121122

122-
namespaceInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
123-
AddFunc: func(obj interface{}) { updateConfigMap() },
124-
UpdateFunc: func(oldObj, newObj interface{}) { updateConfigMap() },
125-
DeleteFunc: func(obj interface{}) { updateConfigMap() },
123+
namespaceInformer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
124+
AddFunc: func(obj interface{}) { updateConfigMap(ctx) },
125+
UpdateFunc: func(oldObj, newObj interface{}) { updateConfigMap(ctx) },
126+
DeleteFunc: func(obj interface{}) { updateConfigMap(ctx) },
126127
})
127128

128-
informerFactory.Start(stopCh)
129-
informerFactory.WaitForCacheSync(stopCh)
129+
informerFactory.Start(ctx.Done())
130+
if !cache.WaitForCacheSync(ctx.Done(), namespaceInformer.Informer().HasSynced) {
131+
return fmt.Errorf("failed to sync namespace informer cache")
132+
}
133+
return nil
130134
}
131135

132136
// updateConfigMap updates the ConfigMap with the latest namespace list
133-
func updateConfigMap() {
137+
func updateConfigMap(ctx context.Context) {
134138
lock.Lock()
135139
defer lock.Unlock()
136140

137141
namespaces, err := getFilteredNamespaces()
138142
if err != nil {
139-
fmt.Printf("Error getting namespaces: %v\n", err)
143+
log.Printf("Error getting namespaces: %v", err)
140144
return
141145
}
142146

143147
newValue := strings.Join(namespaces, ",")
144148

145149
// Fetch the existing ConfigMap
146-
cm, err := clientset.CoreV1().ConfigMaps(configMapNamespace).Get(context.TODO(), configMapName, metav1.GetOptions{})
150+
cm, err := clientset.CoreV1().ConfigMaps(configMapNamespace).Get(ctx, configMapName, metav1.GetOptions{})
147151
if err != nil {
148152
if apierrors.IsNotFound(err) {
149153
// If ConfigMap doesn't exist, create it
150-
createConfigMap(newValue)
154+
createConfigMap(ctx, newValue)
151155
} else {
152-
fmt.Printf("Error getting ConfigMap: %v\n", err)
156+
log.Printf("Error getting ConfigMap: %v", err)
153157
}
154158
return
155159
}
@@ -164,34 +168,40 @@ func updateConfigMap() {
164168
cm.Data = make(map[string]string)
165169
}
166170
cm.Data[watchKey] = newValue
167-
_, err = clientset.CoreV1().ConfigMaps(configMapNamespace).Update(context.TODO(), cm, metav1.UpdateOptions{})
168-
if err != nil {
169-
fmt.Printf("Error updating ConfigMap: %v\n", err)
171+
if _, err := clientset.CoreV1().ConfigMaps(configMapNamespace).Update(ctx, cm, metav1.UpdateOptions{}); err != nil {
172+
log.Printf("Error updating ConfigMap: %v", err)
170173
return
171174
}
172175

173176
// Trigger restart of the main container
174177
restartMainContainer()
175178
}
176179

177-
// getFilteredNamespaces retrieves namespaces with the specified label
180+
// getFilteredNamespaces retrieves namespaces with the specified label from the
181+
// informer cache. The result is sorted so the joined value is stable and does
182+
// not trigger spurious ConfigMap updates (and container restarts) caused by
183+
// non-deterministic cache iteration order.
178184
func getFilteredNamespaces() ([]string, error) {
179-
namespaces, err := clientset.CoreV1().Namespaces().List(context.TODO(), metav1.ListOptions{
180-
LabelSelector: labelSelector,
181-
})
185+
selector, err := labels.Parse(labelSelector)
186+
if err != nil {
187+
return nil, fmt.Errorf("invalid label selector %q: %w", labelSelector, err)
188+
}
189+
190+
namespaces, err := namespaceLister.List(selector)
182191
if err != nil {
183192
return nil, err
184193
}
185194

186-
var nsList []string
187-
for _, ns := range namespaces.Items {
195+
nsList := make([]string, 0, len(namespaces))
196+
for _, ns := range namespaces {
188197
nsList = append(nsList, ns.Name)
189198
}
199+
sort.Strings(nsList)
190200
return nsList, nil
191201
}
192202

193203
// createConfigMap creates a new ConfigMap
194-
func createConfigMap(value string) {
204+
func createConfigMap(ctx context.Context, value string) {
195205
cm := &v1.ConfigMap{
196206
ObjectMeta: metav1.ObjectMeta{
197207
Name: configMapName,
@@ -202,27 +212,25 @@ func createConfigMap(value string) {
202212
},
203213
}
204214

205-
_, err := clientset.CoreV1().ConfigMaps(configMapNamespace).Create(context.TODO(), cm, metav1.CreateOptions{})
206-
if err != nil {
207-
fmt.Printf("Error creating ConfigMap: %v\n", err)
215+
if _, err := clientset.CoreV1().ConfigMaps(configMapNamespace).Create(ctx, cm, metav1.CreateOptions{}); err != nil {
216+
log.Printf("Error creating ConfigMap: %v", err)
208217
}
209218
}
210219

211220
// killMainContainer kills the main container to reload the ConfigMap
212221
func killMainContainer() {
213-
fmt.Println("Restarting main container...")
222+
log.Println("Restarting main container...")
214223

215224
// Find the main container's PID (assumes PID namespace is shared)
216225
pid, err := getMainContainerPID()
217226
if err != nil {
218-
fmt.Printf("Error getting main container PID: %v\n", err)
227+
log.Printf("Error getting main container PID: %v", err)
219228
return
220229
}
221230

222231
// Send SIGTERM to the main container
223-
err = exec.Command("kill", "-SIGTERM", pid).Run()
224-
if err != nil {
225-
fmt.Printf("Error sending SIGTERM to main container: %v\n", err)
232+
if err := exec.Command("kill", "-SIGTERM", pid).Run(); err != nil {
233+
log.Printf("Error sending SIGTERM to main container: %v", err)
226234
}
227235
}
228236

‎main_test.go‎

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@ import (
77
v1 "k8s.io/api/core/v1"
88
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
99
"k8s.io/apimachinery/pkg/runtime"
10+
"k8s.io/client-go/informers"
1011
"k8s.io/client-go/kubernetes/fake"
1112
)
1213

1314
// setupTest configures the package globals for a test run and returns the fake
14-
// clientset so assertions can be made against it. The main-container restart is
15-
// stubbed out so tests never shell out to pgrep/kill.
15+
// clientset so assertions can be made against it. A namespace lister backed by
16+
// the fake clientset is wired up the same way production does, and the
17+
// main-container restart is stubbed out so tests never shell out to pgrep/kill.
1618
func setupTest(t *testing.T, objects ...runtime.Object) *fake.Clientset {
1719
t.Helper()
1820

@@ -21,36 +23,49 @@ func setupTest(t *testing.T, objects ...runtime.Object) *fake.Clientset {
2123
configMapName = "watch-namespace-config"
2224
configMapNamespace = "vault"
2325
watchKey = "WATCH_NAMESPACE"
24-
// Empty selector: the fake clientset does not apply label-selector filtering,
25-
// so keep tests independent of that behaviour.
26+
// Empty selector matches every namespace; the lister applies the selector
27+
// in-memory, so tests stay independent of any pre-filtering.
2628
labelSelector = ""
2729

2830
clientset = cs
2931

32+
// Build and sync a namespace lister from the fake clientset.
33+
factory := informers.NewSharedInformerFactory(cs, 0)
34+
namespaceLister = factory.Core().V1().Namespaces().Lister()
35+
stopCh := make(chan struct{})
36+
factory.Start(stopCh)
37+
factory.WaitForCacheSync(stopCh)
38+
3039
restartMainContainer = func() {}
31-
t.Cleanup(func() { restartMainContainer = killMainContainer })
40+
t.Cleanup(func() {
41+
close(stopCh)
42+
restartMainContainer = killMainContainer
43+
namespaceLister = nil
44+
})
3245

3346
return cs
3447
}
3548

3649
func TestGetFilteredNamespaces(t *testing.T) {
3750
setupTest(t,
38-
&v1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "team-a"}},
3951
&v1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "team-b"}},
52+
&v1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "team-a"}},
4053
)
4154

4255
got, err := getFilteredNamespaces()
4356
if err != nil {
4457
t.Fatalf("getFilteredNamespaces returned error: %v", err)
4558
}
4659

47-
want := map[string]bool{"team-a": true, "team-b": true}
60+
// Result must be sorted for a stable, restart-friendly ConfigMap value.
61+
want := []string{"team-a", "team-b"}
4862
if len(got) != len(want) {
49-
t.Fatalf("got %d namespaces (%v), want %d", len(got), got, len(want))
63+
t.Fatalf("got %d namespaces (%v), want %v", len(got), got, want)
5064
}
51-
for _, ns := range got {
52-
if !want[ns] {
53-
t.Errorf("unexpected namespace %q in result", ns)
65+
for i := range want {
66+
if got[i] != want[i] {
67+
t.Errorf("got %v, want %v (not sorted)", got, want)
68+
break
5469
}
5570
}
5671
}
@@ -60,7 +75,7 @@ func TestUpdateConfigMap_CreatesWhenMissing(t *testing.T) {
6075
&v1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "team-a"}},
6176
)
6277

63-
updateConfigMap()
78+
updateConfigMap(context.Background())
6479

6580
cm, err := cs.CoreV1().ConfigMaps(configMapNamespace).Get(context.TODO(), configMapName, metav1.GetOptions{})
6681
if err != nil {
@@ -83,7 +98,7 @@ func TestUpdateConfigMap_NilData(t *testing.T) {
8398
)
8499

85100
// Would panic before the nil-map guard was added.
86-
updateConfigMap()
101+
updateConfigMap(context.Background())
87102

88103
cm, err := cs.CoreV1().ConfigMaps(configMapNamespace).Get(context.TODO(), configMapName, metav1.GetOptions{})
89104
if err != nil {
@@ -103,7 +118,7 @@ func TestUpdateConfigMap_NoChangeKeepsValue(t *testing.T) {
103118
},
104119
)
105120

106-
updateConfigMap()
121+
updateConfigMap(context.Background())
107122

108123
cm, err := cs.CoreV1().ConfigMaps(configMapNamespace).Get(context.TODO(), configMapName, metav1.GetOptions{})
109124
if err != nil {

0 commit comments

Comments
 (0)