Memory ballooning support on firecracker-containerd

Implement the memory ballooning APIs on firecracker-containerd

Signed-off-by: Royce Zhao <[email protected]>

Author: RoyceDavison

firecracker-control/local.go

@@ -377,6 +377,78 @@ func (s *local) GetVMMetadata(requestCtx context.Context, req *proto.GetVMMetada
 	return resp, nil
 }
 
+// GetBalloonConfig get balloon configuration, before or after machine startup.
+func (s *local) GetBalloonConfig(requestCtx context.Context, req *proto.GetBalloonConfigRequest) (*proto.GetBalloonConfigResponse, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
+	}
+
+	defer client.Close()
+	resp, err := client.GetBalloonConfig(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to get balloon config")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return resp, nil
+}
+
+// UpdateBalloon updates memory size for an existing balloon device, before or after machine startup.
+func (s *local) UpdateBalloon(requestCtx context.Context, req *proto.UpdateBalloonRequest) (*empty.Empty, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
+	}
+
+	defer client.Close()
+	resp, err := client.UpdateBalloon(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to update balloon")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return resp, nil
+}
+
+// GetBalloonStats will return the latest balloon device statistics, only if enabled pre-boot.
+func (s *local) GetBalloonStats(requestCtx context.Context, req *proto.GetBalloonStatsRequest) (*proto.GetBalloonStatsResponse, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
+	}
+
+	defer client.Close()
+	resp, err := client.GetBalloonStats(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to get balloon statistics")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return resp, nil
+}
+
+// UpdateBalloonStats updates an existing balloon device statistics interval, before or after machine startup.
+func (s *local) UpdateBalloonStats(requestCtx context.Context, req *proto.UpdateBalloonStatsRequest) (*empty.Empty, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
+	}
+
+	defer client.Close()
+	resp, err := client.UpdateBalloonStats(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to update balloon interval")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return resp, nil
+}
+
 func (s *local) newShim(ns, vmID, containerdAddress string, shimSocket *net.UnixListener, fcSocket *net.UnixListener) (*exec.Cmd, error) {
 	logger := s.logger.WithField("vmID", vmID)
 

firecracker-control/service.go

@@ -106,3 +106,23 @@ func (s *service) GetVMMetadata(ctx context.Context, req *proto.GetVMMetadataReq
 	log.G(ctx).Debug("Getting vm metadata")
 	return s.local.GetVMMetadata(ctx, req)
 }
+
+func (s *service) GetBalloonConfig(ctx context.Context, req *proto.GetBalloonConfigRequest) (*proto.GetBalloonConfigResponse, error) {
+	log.G(ctx).Debug("Getting balloon configuration")
+	return s.local.GetBalloonConfig(ctx, req)
+}
+
+func (s *service) UpdateBalloon(ctx context.Context, req *proto.UpdateBalloonRequest) (*empty.Empty, error) {
+	log.G(ctx).Debug("Updating balloon memory size")
+	return s.local.UpdateBalloon(ctx, req)
+}
+
+func (s *service) GetBalloonStats(ctx context.Context, req *proto.GetBalloonStatsRequest) (*proto.GetBalloonStatsResponse, error) {
+	log.G(ctx).Debug("Getting balloon statistics")
+	return s.local.GetBalloonStats(ctx, req)
+}
+
+func (s *service) UpdateBalloonStats(ctx context.Context, req *proto.UpdateBalloonStatsRequest) (*empty.Empty, error) {
+	log.G(ctx).Debug("Updating balloon device statistics polling interval")
+	return s.local.UpdateBalloonStats(ctx, req)
+}