Skip to content

Commit b2ea3a4

Browse files
drewmalindrewmalin
committed
save
1 parent fe3cf64 commit b2ea3a4

1 file changed

Lines changed: 143 additions & 146 deletions

File tree

‎v1/providers/testkube/instance.go‎

Lines changed: 143 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -65,65 +65,152 @@ func (c *TestKubeClient) CreateInstance(ctx context.Context, attrs cloudv1.Creat
6565
return nil, cloudv1.ErrOutOfQuota
6666
}
6767

68+
instance, err := c.createInstanceAsK8sResources(ctx, attrs, instanceTypeSpec)
69+
if err != nil {
70+
return nil, err
71+
}
72+
return instance, nil
73+
}
74+
75+
func (c *TestKubeClient) createInstanceAsK8sResources(ctx context.Context, attrs cloudv1.CreateInstanceAttrs, instanceTypeSpec instanceTypeSpec) (*cloudv1.Instance, error) {
6876
// Create a "cloud ID" to emulate a provider-provided instance ID.
6977
cloudID := makeCloudID(c.refID, attrs.RefID)
7078

71-
if _, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Get(ctx, string(cloudID), metav1.GetOptions{}); err == nil {
72-
return c.GetInstance(ctx, cloudID)
73-
} else if !apierrors.IsNotFound(err) {
74-
return nil, fmt.Errorf("get existing testkube instance: %w", err)
75-
}
76-
77-
service := c.newInstanceAsK8sService(cloudID, attrs, instanceTypeSpec)
78-
serviceCreated := false
79-
createdService, err := c.k8sClient.CoreV1().Services(c.namespace).Create(ctx, service, metav1.CreateOptions{})
79+
location := c.resourceLocation(attrs)
80+
annotations := c.resourceAnnotations(cloudID, attrs, instanceTypeSpec)
81+
82+
// Create the service.
83+
k8sService, err := c.k8sClient.
84+
CoreV1().
85+
Services(c.namespace).
86+
Create(ctx, &corev1.Service{
87+
ObjectMeta: metav1.ObjectMeta{
88+
Name: string(cloudID),
89+
Namespace: c.namespace,
90+
Labels: objectLabels(string(cloudID), location),
91+
Annotations: annotations,
92+
},
93+
Spec: corev1.ServiceSpec{
94+
Type: instanceTypeSpec.serviceType,
95+
Selector: selectorLabels(string(cloudID)),
96+
Ports: []corev1.ServicePort{
97+
{
98+
Name: servicePortName,
99+
Protocol: corev1.ProtocolTCP,
100+
Port: servicePort,
101+
TargetPort: intstr.FromInt32(containerSSHPort),
102+
},
103+
},
104+
},
105+
}, metav1.CreateOptions{})
80106
if err != nil {
81-
if !apierrors.IsAlreadyExists(err) {
82-
return nil, fmt.Errorf("create testkube service: %w", err)
83-
}
84-
createdService, err = c.k8sClient.CoreV1().Services(c.namespace).Get(ctx, string(cloudID), metav1.GetOptions{})
85-
if err != nil {
86-
return nil, fmt.Errorf("get existing testkube service: %w", err)
87-
}
88-
} else {
89-
serviceCreated = true
107+
return nil, fmt.Errorf("create testkube service: %w", err)
90108
}
91109

92-
statefulSet := c.makeStatefulSet(cloudID, attrs, instanceTypeSpec)
93-
createdStatefulSet, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Create(ctx, statefulSet, metav1.CreateOptions{})
110+
// Create the stateful set.
111+
replicas := int32(1)
112+
k8sStatefulSet, err := c.k8sClient.
113+
AppsV1().
114+
StatefulSets(c.namespace).
115+
Create(ctx, &appsv1.StatefulSet{
116+
ObjectMeta: metav1.ObjectMeta{
117+
Name: string(cloudID),
118+
Namespace: c.namespace,
119+
Labels: objectLabels(string(cloudID), location),
120+
Annotations: annotations,
121+
},
122+
Spec: appsv1.StatefulSetSpec{
123+
Replicas: &replicas,
124+
ServiceName: string(cloudID),
125+
Selector: &metav1.LabelSelector{
126+
MatchLabels: selectorLabels(string(cloudID)),
127+
},
128+
Template: corev1.PodTemplateSpec{
129+
ObjectMeta: metav1.ObjectMeta{
130+
Labels: objectLabels(string(cloudID), location),
131+
Annotations: annotations,
132+
},
133+
Spec: corev1.PodSpec{
134+
TerminationGracePeriodSeconds: int64Ptr(1),
135+
Containers: []corev1.Container{
136+
{
137+
Name: "vm",
138+
Image: instanceTypeSpec.image,
139+
ImagePullPolicy: corev1.PullIfNotPresent,
140+
SecurityContext: &corev1.SecurityContext{
141+
Privileged: boolPtr(true),
142+
},
143+
Ports: []corev1.ContainerPort{
144+
{
145+
Name: servicePortName,
146+
ContainerPort: containerSSHPort,
147+
Protocol: corev1.ProtocolTCP,
148+
},
149+
},
150+
ReadinessProbe: &corev1.Probe{
151+
ProbeHandler: corev1.ProbeHandler{
152+
TCPSocket: &corev1.TCPSocketAction{
153+
Port: intstr.FromInt32(containerSSHPort),
154+
},
155+
},
156+
InitialDelaySeconds: 1,
157+
PeriodSeconds: 2,
158+
FailureThreshold: 30,
159+
},
160+
Env: c.containerEnv(attrs),
161+
Resources: corev1.ResourceRequirements{
162+
Requests: corev1.ResourceList{
163+
corev1.ResourceCPU: resource.MustParse("250m"),
164+
corev1.ResourceMemory: resource.MustParse("512Mi"),
165+
},
166+
Limits: corev1.ResourceList{
167+
corev1.ResourceCPU: resource.MustParse("2"),
168+
corev1.ResourceMemory: resource.MustParse("4Gi"),
169+
},
170+
},
171+
},
172+
},
173+
},
174+
},
175+
},
176+
}, metav1.CreateOptions{})
94177
if err != nil {
95-
if apierrors.IsAlreadyExists(err) {
96-
return c.GetInstance(ctx, cloudID)
97-
}
98-
if serviceCreated {
99-
_ = c.k8sClient.CoreV1().Services(c.namespace).Delete(ctx, string(cloudID), metav1.DeleteOptions{})
100-
}
101178
return nil, fmt.Errorf("create testkube statefulset: %w", err)
102179
}
103180

104-
return c.instanceFromResources(createdStatefulSet, createdService, nil), nil
181+
// Map to the brev instance.
182+
return c.instanceFromResources(k8sStatefulSet, k8sService, nil), nil
105183
}
106184

107185
func (c *TestKubeClient) GetInstance(ctx context.Context, instanceID cloudv1.CloudProviderInstanceID) (*cloudv1.Instance, error) {
108-
statefulSet, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Get(ctx, string(instanceID), metav1.GetOptions{})
186+
statefulSet, err := c.k8sClient.
187+
AppsV1().
188+
StatefulSets(c.namespace).
189+
Get(ctx, string(instanceID), metav1.GetOptions{})
109190
if err != nil {
110191
if apierrors.IsNotFound(err) {
111192
return nil, fmt.Errorf("%w: %s", cloudv1.ErrInstanceNotFound, instanceID)
112193
}
113194
return nil, fmt.Errorf("get testkube statefulset: %w", err)
114195
}
115196

116-
service, err := c.k8sClient.CoreV1().Services(c.namespace).Get(ctx, string(instanceID), metav1.GetOptions{})
197+
service, err := c.k8sClient.
198+
CoreV1().
199+
Services(c.namespace).
200+
Get(ctx, string(instanceID), metav1.GetOptions{})
117201
if err != nil && !apierrors.IsNotFound(err) {
118202
return nil, fmt.Errorf("get testkube service: %w", err)
119203
}
120204
if apierrors.IsNotFound(err) {
121205
service = nil
122206
}
123207

124-
pods, err := c.k8sClient.CoreV1().Pods(c.namespace).List(ctx, metav1.ListOptions{
125-
LabelSelector: labels.SelectorFromSet(selectorLabels(string(instanceID))).String(),
126-
})
208+
pods, err := c.k8sClient.
209+
CoreV1().
210+
Pods(c.namespace).
211+
List(ctx, metav1.ListOptions{
212+
LabelSelector: labels.SelectorFromSet(selectorLabels(string(instanceID))).String(),
213+
})
127214
if err != nil {
128215
return nil, fmt.Errorf("list testkube pods: %w", err)
129216
}
@@ -132,12 +219,15 @@ func (c *TestKubeClient) GetInstance(ctx context.Context, instanceID cloudv1.Clo
132219
}
133220

134221
func (c *TestKubeClient) ListInstances(ctx context.Context, args cloudv1.ListInstancesArgs) ([]cloudv1.Instance, error) {
135-
statefulSets, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).List(ctx, metav1.ListOptions{
136-
LabelSelector: labels.SelectorFromSet(labels.Set{
137-
labelManagedBy: labelManagedByValue,
138-
labelName: labelNameValue,
139-
}).String(),
140-
})
222+
statefulSets, err := c.k8sClient.
223+
AppsV1().
224+
StatefulSets(c.namespace).
225+
List(ctx, metav1.ListOptions{
226+
LabelSelector: labels.SelectorFromSet(labels.Set{
227+
labelManagedBy: labelManagedByValue,
228+
labelName: labelNameValue,
229+
}).String(),
230+
})
141231
if err != nil {
142232
return nil, fmt.Errorf("list testkube statefulsets: %w", err)
143233
}
@@ -186,6 +276,21 @@ func (c *TestKubeClient) StartInstance(ctx context.Context, instanceID cloudv1.C
186276
return c.updateReplicas(ctx, instanceID, 1)
187277
}
188278

279+
func (c *TestKubeClient) updateReplicas(ctx context.Context, instanceID cloudv1.CloudProviderInstanceID, replicas int32) error {
280+
statefulSet, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Get(ctx, string(instanceID), metav1.GetOptions{})
281+
if err != nil {
282+
if apierrors.IsNotFound(err) {
283+
return fmt.Errorf("%w: %s", cloudv1.ErrInstanceNotFound, instanceID)
284+
}
285+
return fmt.Errorf("get testkube statefulset: %w", err)
286+
}
287+
statefulSet.Spec.Replicas = int32Ptr(replicas)
288+
if _, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Update(ctx, statefulSet, metav1.UpdateOptions{}); err != nil {
289+
return fmt.Errorf("update testkube replicas: %w", err)
290+
}
291+
return nil
292+
}
293+
189294
func (c *TestKubeClient) RebootInstance(ctx context.Context, instanceID cloudv1.CloudProviderInstanceID) error {
190295
if _, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Get(ctx, string(instanceID), metav1.GetOptions{}); err != nil {
191296
if apierrors.IsNotFound(err) {
@@ -254,114 +359,6 @@ func (c *TestKubeClient) MergeInstanceTypeForUpdate(_ cloudv1.InstanceType, newI
254359
return newIt
255360
}
256361

257-
func (c *TestKubeClient) updateReplicas(ctx context.Context, instanceID cloudv1.CloudProviderInstanceID, replicas int32) error {
258-
statefulSet, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Get(ctx, string(instanceID), metav1.GetOptions{})
259-
if err != nil {
260-
if apierrors.IsNotFound(err) {
261-
return fmt.Errorf("%w: %s", cloudv1.ErrInstanceNotFound, instanceID)
262-
}
263-
return fmt.Errorf("get testkube statefulset: %w", err)
264-
}
265-
statefulSet.Spec.Replicas = int32Ptr(replicas)
266-
if _, err := c.k8sClient.AppsV1().StatefulSets(c.namespace).Update(ctx, statefulSet, metav1.UpdateOptions{}); err != nil {
267-
return fmt.Errorf("update testkube replicas: %w", err)
268-
}
269-
return nil
270-
}
271-
272-
func (c *TestKubeClient) newInstanceAsK8sService(cloudID cloudv1.CloudProviderInstanceID, attrs cloudv1.CreateInstanceAttrs, spec instanceTypeSpec) *corev1.Service {
273-
location := c.resourceLocation(attrs)
274-
return &corev1.Service{
275-
ObjectMeta: metav1.ObjectMeta{
276-
Name: string(cloudID),
277-
Namespace: c.namespace,
278-
Labels: objectLabels(string(cloudID), location),
279-
Annotations: c.resourceAnnotations(cloudID, attrs, spec),
280-
},
281-
Spec: corev1.ServiceSpec{
282-
Type: spec.serviceType,
283-
Selector: selectorLabels(string(cloudID)),
284-
Ports: []corev1.ServicePort{
285-
{
286-
Name: servicePortName,
287-
Protocol: corev1.ProtocolTCP,
288-
Port: servicePort,
289-
TargetPort: intstr.FromInt32(containerSSHPort),
290-
},
291-
},
292-
},
293-
}
294-
}
295-
296-
func (c *TestKubeClient) makeStatefulSet(cloudID cloudv1.CloudProviderInstanceID, attrs cloudv1.CreateInstanceAttrs, spec instanceTypeSpec) *appsv1.StatefulSet {
297-
replicas := int32(1)
298-
annotations := c.resourceAnnotations(cloudID, attrs, spec)
299-
location := c.resourceLocation(attrs)
300-
podSpec := corev1.PodSpec{
301-
TerminationGracePeriodSeconds: int64Ptr(1),
302-
Containers: []corev1.Container{
303-
{
304-
Name: "vm",
305-
Image: spec.image,
306-
ImagePullPolicy: corev1.PullIfNotPresent,
307-
SecurityContext: &corev1.SecurityContext{
308-
Privileged: boolPtr(true),
309-
},
310-
Ports: []corev1.ContainerPort{
311-
{
312-
Name: servicePortName,
313-
ContainerPort: containerSSHPort,
314-
Protocol: corev1.ProtocolTCP,
315-
},
316-
},
317-
ReadinessProbe: &corev1.Probe{
318-
ProbeHandler: corev1.ProbeHandler{
319-
TCPSocket: &corev1.TCPSocketAction{
320-
Port: intstr.FromInt32(containerSSHPort),
321-
},
322-
},
323-
InitialDelaySeconds: 1,
324-
PeriodSeconds: 2,
325-
FailureThreshold: 30,
326-
},
327-
Env: c.containerEnv(attrs),
328-
Resources: corev1.ResourceRequirements{
329-
Requests: corev1.ResourceList{
330-
corev1.ResourceCPU: resource.MustParse("250m"),
331-
corev1.ResourceMemory: resource.MustParse("512Mi"),
332-
},
333-
Limits: corev1.ResourceList{
334-
corev1.ResourceCPU: resource.MustParse("2"),
335-
corev1.ResourceMemory: resource.MustParse("4Gi"),
336-
},
337-
},
338-
},
339-
},
340-
}
341-
return &appsv1.StatefulSet{
342-
ObjectMeta: metav1.ObjectMeta{
343-
Name: string(cloudID),
344-
Namespace: c.namespace,
345-
Labels: objectLabels(string(cloudID), location),
346-
Annotations: annotations,
347-
},
348-
Spec: appsv1.StatefulSetSpec{
349-
Replicas: &replicas,
350-
ServiceName: string(cloudID),
351-
Selector: &metav1.LabelSelector{
352-
MatchLabels: selectorLabels(string(cloudID)),
353-
},
354-
Template: corev1.PodTemplateSpec{
355-
ObjectMeta: metav1.ObjectMeta{
356-
Labels: objectLabels(string(cloudID), location),
357-
Annotations: annotations,
358-
},
359-
Spec: podSpec,
360-
},
361-
},
362-
}
363-
}
364-
365362
func (c *TestKubeClient) resourceAnnotations(cloudID cloudv1.CloudProviderInstanceID, attrs cloudv1.CreateInstanceAttrs, spec instanceTypeSpec) map[string]string {
366363
name := attrs.Name
367364
if name == "" {

0 commit comments

Comments
 (0)