xref: /openbmc/qemu/hw/hyperv/hv-balloon.c (revision ecd6f6a8)
1 /*
2  * QEMU Hyper-V Dynamic Memory Protocol driver
3  *
4  * Copyright (C) 2020-2023 Oracle and/or its affiliates.
5  *
6  * This work is licensed under the terms of the GNU GPL, version 2 or later.
7  * See the COPYING file in the top-level directory.
8  */
9 
10 #include "qemu/osdep.h"
11 #include "hv-balloon-internal.h"
12 
13 #include "exec/address-spaces.h"
14 #include "exec/cpu-common.h"
15 #include "exec/ramblock.h"
16 #include "hw/boards.h"
17 #include "hw/hyperv/dynmem-proto.h"
18 #include "hw/hyperv/hv-balloon.h"
19 #include "hw/hyperv/vmbus.h"
20 #include "hw/mem/memory-device.h"
21 #include "hw/mem/pc-dimm.h"
22 #include "hw/qdev-core.h"
23 #include "hw/qdev-properties.h"
24 #include "monitor/qdev.h"
25 #include "qapi/error.h"
26 #include "qapi/qapi-commands-machine.h"
27 #include "qapi/qapi-events-machine.h"
28 #include "qapi/qapi-types-machine.h"
29 #include "qapi/qmp/qdict.h"
30 #include "qapi/visitor.h"
31 #include "qemu/error-report.h"
32 #include "qemu/module.h"
33 #include "qemu/units.h"
34 #include "qemu/timer.h"
35 #include "sysemu/balloon.h"
36 #include "sysemu/hostmem.h"
37 #include "sysemu/reset.h"
38 #include "hv-balloon-our_range_memslots.h"
39 #include "hv-balloon-page_range_tree.h"
40 #include "trace.h"
41 
42 #define HV_BALLOON_ADDR_PROP "addr"
43 #define HV_BALLOON_MEMDEV_PROP "memdev"
44 #define HV_BALLOON_GUID "525074DC-8985-46e2-8057-A307DC18A502"
45 
46 /*
47  * Some Windows versions (at least Server 2019) will crash with various
48  * error codes when receiving DM protocol requests (at least
49  * DM_MEM_HOT_ADD_REQUEST) immediately after boot.
50  *
51  * It looks like Hyper-V from Server 2016 uses a 50-second after-boot
52  * delay, probably to workaround this issue, so we'll use this value, too.
53  */
54 #define HV_BALLOON_POST_INIT_WAIT (50 * 1000)
55 
56 #define HV_BALLOON_HA_CHUNK_SIZE (2 * GiB)
57 #define HV_BALLOON_HA_CHUNK_PAGES (HV_BALLOON_HA_CHUNK_SIZE / HV_BALLOON_PAGE_SIZE)
58 
59 #define HV_BALLOON_HA_MEMSLOT_SIZE_ALIGN (128 * MiB)
60 
61 #define HV_BALLOON_HR_CHUNK_PAGES 585728
62 /*
63  *                                ^ that's the maximum number of pages
64  * that Windows returns in one hot remove response
65  *
66  * If the number requested is too high Windows will no longer honor
67  * these requests
68  */
69 
70 struct HvBalloonClass {
71     VMBusDeviceClass parent_class;
72 } HvBalloonClass;
73 
74 typedef enum State {
75     /* not a real state */
76     S_NO_CHANGE = 0,
77 
78     S_WAIT_RESET,
79     S_POST_RESET_CLOSED,
80 
81     /* init flow */
82     S_VERSION,
83     S_CAPS,
84     S_POST_INIT_WAIT,
85 
86     S_IDLE,
87 
88     /* balloon op flow */
89     S_BALLOON_POSTING,
90     S_BALLOON_RB_WAIT,
91     S_BALLOON_REPLY_WAIT,
92 
93     /* unballoon + hot add ops flow */
94     S_UNBALLOON_POSTING,
95     S_UNBALLOON_RB_WAIT,
96     S_UNBALLOON_REPLY_WAIT,
97     S_HOT_ADD_SETUP,
98     S_HOT_ADD_RB_WAIT,
99     S_HOT_ADD_POSTING,
100     S_HOT_ADD_REPLY_WAIT,
101 } State;
102 
103 typedef struct StateDesc {
104     State state;
105     const char *desc;
106 } StateDesc;
107 
108 typedef struct HvBalloon {
109     VMBusDevice parent;
110     State state;
111 
112     union dm_version version;
113     union dm_caps caps;
114 
115     QEMUTimer post_init_timer;
116 
117     unsigned int trans_id;
118 
119     struct {
120         bool enabled;
121         bool received;
122         uint64_t committed;
123         uint64_t available;
124     } status_report;
125 
126     /* Guest target size */
127     uint64_t target;
128     bool target_changed;
129 
130     /* Current (un)balloon / hot-add operation parameters */
131     union {
132         uint64_t balloon_diff;
133 
134         struct {
135             uint64_t unballoon_diff;
136             uint64_t hot_add_diff;
137         };
138 
139         struct {
140             PageRange hot_add_range;
141             uint64_t ha_current_count;
142         };
143     };
144 
145     OurRangeMemslots *our_range;
146 
147     /* Count of memslots covering our memory */
148     unsigned int memslot_count;
149 
150     /* Nominal size of each memslot (the last one might be smaller) */
151     uint64_t memslot_size;
152 
153     /* Non-ours removed memory */
154     PageRangeTree removed_guest, removed_both;
155 
156     /* Grand totals of removed memory (both ours and non-ours) */
157     uint64_t removed_guest_ctr, removed_both_ctr;
158 
159     /* MEMORY_DEVICE props */
160     uint64_t addr;
161     HostMemoryBackend *hostmem;
162     MemoryRegion *mr;
163 } HvBalloon;
164 
165 OBJECT_DEFINE_TYPE_WITH_INTERFACES(HvBalloon, hv_balloon, HV_BALLOON, VMBUS_DEVICE, \
166                                    { TYPE_MEMORY_DEVICE }, { })
167 
168 #define HV_BALLOON_SET_STATE(hvb, news)             \
169     do {                                            \
170         assert(news != S_NO_CHANGE);                \
171         hv_balloon_state_set(hvb, news, # news);    \
172     } while (0)
173 
174 #define HV_BALLOON_STATE_DESC_SET(stdesc, news)         \
175     _hv_balloon_state_desc_set(stdesc, news, # news)
176 
177 #define HV_BALLOON_STATE_DESC_INIT \
178     {                              \
179         .state = S_NO_CHANGE,      \
180     }
181 
182 typedef struct HvBalloonReq {
183     VMBusChanReq vmreq;
184 } HvBalloonReq;
185 
186 /* total our memory includes parts currently removed from the guest */
187 static uint64_t hv_balloon_total_our_ram(HvBalloon *balloon)
188 {
189     if (!balloon->our_range) {
190         return 0;
191     }
192 
193     return balloon->our_range->range.added;
194 }
195 
196 /* TODO: unify the code below with virtio-balloon and cache the value */
197 static int build_dimm_list(Object *obj, void *opaque)
198 {
199     GSList **list = opaque;
200 
201     if (object_dynamic_cast(obj, TYPE_PC_DIMM)) {
202         DeviceState *dev = DEVICE(obj);
203         if (dev->realized) { /* only realized DIMMs matter */
204             *list = g_slist_prepend(*list, dev);
205         }
206     }
207 
208     object_child_foreach(obj, build_dimm_list, opaque);
209     return 0;
210 }
211 
212 static ram_addr_t get_current_ram_size(void)
213 {
214     GSList *list = NULL, *item;
215     ram_addr_t size = current_machine->ram_size;
216 
217     build_dimm_list(qdev_get_machine(), &list);
218     for (item = list; item; item = g_slist_next(item)) {
219         Object *obj = OBJECT(item->data);
220         if (!strcmp(object_get_typename(obj), TYPE_PC_DIMM))
221             size += object_property_get_int(obj, PC_DIMM_SIZE_PROP,
222                                             &error_abort);
223     }
224     g_slist_free(list);
225 
226     return size;
227 }
228 
229 /* total RAM includes memory currently removed from the guest */
230 static uint64_t hv_balloon_total_ram(HvBalloon *balloon)
231 {
232     ram_addr_t ram_size = get_current_ram_size();
233     uint64_t ram_size_pages = ram_size >> HV_BALLOON_PFN_SHIFT;
234     uint64_t our_ram_size_pages = hv_balloon_total_our_ram(balloon);
235 
236     assert(ram_size_pages > 0);
237 
238     return SUM_SATURATE_U64(ram_size_pages, our_ram_size_pages);
239 }
240 
241 /*
242  * calculating the total RAM size is a slow operation,
243  * avoid it as much as possible
244  */
245 static uint64_t hv_balloon_total_removed_rs(HvBalloon *balloon,
246                                             uint64_t ram_size_pages)
247 {
248     uint64_t total_removed;
249 
250     total_removed = SUM_SATURATE_U64(balloon->removed_guest_ctr,
251                                      balloon->removed_both_ctr);
252 
253     /* possible if guest returns pages outside actual RAM */
254     if (total_removed > ram_size_pages) {
255         total_removed = ram_size_pages;
256     }
257 
258     return total_removed;
259 }
260 
261 /* Returns whether the state has actually changed */
262 static bool hv_balloon_state_set(HvBalloon *balloon,
263                                  State newst, const char *newststr)
264 {
265     if (newst == S_NO_CHANGE || balloon->state == newst) {
266         return false;
267     }
268 
269     balloon->state = newst;
270     trace_hv_balloon_state_change(newststr);
271     return true;
272 }
273 
274 static void _hv_balloon_state_desc_set(StateDesc *stdesc,
275                                        State newst, const char *newststr)
276 {
277     /* state setting is only permitted on a freshly init desc */
278     assert(stdesc->state == S_NO_CHANGE);
279 
280     assert(newst != S_NO_CHANGE);
281 
282     stdesc->state = newst;
283     stdesc->desc = newststr;
284 }
285 
286 static VMBusChannel *hv_balloon_get_channel_maybe(HvBalloon *balloon)
287 {
288     return vmbus_device_channel(&balloon->parent, 0);
289 }
290 
291 static VMBusChannel *hv_balloon_get_channel(HvBalloon *balloon)
292 {
293     VMBusChannel *chan;
294 
295     chan = hv_balloon_get_channel_maybe(balloon);
296     assert(chan != NULL);
297     return chan;
298 }
299 
300 static ssize_t hv_balloon_send_packet(VMBusChannel *chan,
301                                       struct dm_message *msg)
302 {
303     int ret;
304 
305     ret = vmbus_channel_reserve(chan, 0, msg->hdr.size);
306     if (ret < 0) {
307         return ret;
308     }
309 
310     return vmbus_channel_send(chan, VMBUS_PACKET_DATA_INBAND,
311                               NULL, 0, msg, msg->hdr.size, false,
312                               msg->hdr.trans_id);
313 }
314 
315 static bool hv_balloon_unballoon_get_source(HvBalloon *balloon,
316                                             PageRangeTree *dtree,
317                                             uint64_t **dctr,
318                                             bool *is_our_range)
319 {
320     OurRange *our_range = OUR_RANGE(balloon->our_range);
321 
322     /* Try the boot memory first */
323     if (g_tree_nnodes(balloon->removed_guest.t) > 0) {
324         *dtree = balloon->removed_guest;
325         *dctr = &balloon->removed_guest_ctr;
326         *is_our_range = false;
327     } else if (g_tree_nnodes(balloon->removed_both.t) > 0) {
328         *dtree = balloon->removed_both;
329         *dctr = &balloon->removed_both_ctr;
330         *is_our_range = false;
331     } else if (!our_range) {
332         return false;
333     } else if (!our_range_is_removed_tree_empty(our_range, false)) {
334         *dtree = our_range_get_removed_tree(our_range, false);
335         *dctr = &balloon->removed_guest_ctr;
336         *is_our_range = true;
337     } else if (!our_range_is_removed_tree_empty(our_range, true)) {
338         *dtree = our_range_get_removed_tree(our_range, true);
339         *dctr = &balloon->removed_both_ctr;
340         *is_our_range = true;
341     } else {
342         return false;
343     }
344 
345     return true;
346 }
347 
348 static void hv_balloon_unballoon_rb_wait(HvBalloon *balloon, StateDesc *stdesc)
349 {
350     VMBusChannel *chan = hv_balloon_get_channel(balloon);
351     struct dm_unballoon_request *ur;
352     size_t ur_size = sizeof(*ur) + sizeof(ur->range_array[0]);
353 
354     assert(balloon->state == S_UNBALLOON_RB_WAIT);
355 
356     if (vmbus_channel_reserve(chan, 0, ur_size) < 0) {
357         return;
358     }
359 
360     HV_BALLOON_STATE_DESC_SET(stdesc, S_UNBALLOON_POSTING);
361 }
362 
363 static void hv_balloon_unballoon_posting(HvBalloon *balloon, StateDesc *stdesc)
364 {
365     VMBusChannel *chan = hv_balloon_get_channel(balloon);
366     PageRangeTree dtree;
367     uint64_t *dctr;
368     bool our_range;
369     struct dm_unballoon_request *ur;
370     size_t ur_size = sizeof(*ur) + sizeof(ur->range_array[0]);
371     PageRange range;
372     bool bret;
373     ssize_t ret;
374 
375     assert(balloon->state == S_UNBALLOON_POSTING);
376     assert(balloon->unballoon_diff > 0);
377 
378     if (!hv_balloon_unballoon_get_source(balloon, &dtree, &dctr, &our_range)) {
379         error_report("trying to unballoon but nothing seems to be ballooned");
380         /*
381          * there is little we can do as we might have already
382          * sent the guest a partial request we can't cancel
383          */
384         return;
385     }
386 
387     assert(balloon->our_range || !our_range);
388     assert(dtree.t);
389     assert(dctr);
390 
391     ur = alloca(ur_size);
392     memset(ur, 0, ur_size);
393     ur->hdr.type = DM_UNBALLOON_REQUEST;
394     ur->hdr.size = ur_size;
395     ur->hdr.trans_id = balloon->trans_id;
396 
397     bret = hvb_page_range_tree_pop(dtree, &range, MIN(balloon->unballoon_diff,
398                                                       HV_BALLOON_HA_CHUNK_PAGES));
399     assert(bret);
400     /* TODO: madvise? */
401 
402     *dctr -= range.count;
403     balloon->unballoon_diff -= range.count;
404 
405     ur->range_count = 1;
406     ur->range_array[0].finfo.start_page = range.start;
407     ur->range_array[0].finfo.page_cnt = range.count;
408     ur->more_pages = balloon->unballoon_diff > 0;
409 
410     trace_hv_balloon_outgoing_unballoon(ur->hdr.trans_id,
411                                         range.count, range.start,
412                                         balloon->unballoon_diff);
413 
414     if (ur->more_pages) {
415         HV_BALLOON_STATE_DESC_SET(stdesc, S_UNBALLOON_RB_WAIT);
416     } else {
417         HV_BALLOON_STATE_DESC_SET(stdesc, S_UNBALLOON_REPLY_WAIT);
418     }
419 
420     ret = vmbus_channel_send(chan, VMBUS_PACKET_DATA_INBAND,
421                              NULL, 0, ur, ur_size, false,
422                              ur->hdr.trans_id);
423     if (ret <= 0) {
424         error_report("error %zd when posting unballoon msg, expect problems",
425                      ret);
426     }
427 }
428 
429 static bool hv_balloon_our_range_ensure(HvBalloon *balloon)
430 {
431     uint64_t align;
432     MemoryRegion *hostmem_mr;
433     g_autoptr(OurRangeMemslots) our_range_memslots = NULL;
434     OurRange *our_range;
435 
436     if (balloon->our_range) {
437         return true;
438     }
439 
440     if (!balloon->hostmem) {
441         return false;
442     }
443 
444     align = (1 << balloon->caps.cap_bits.hot_add_alignment) * MiB;
445     assert(QEMU_IS_ALIGNED(balloon->addr, align));
446 
447     hostmem_mr = host_memory_backend_get_memory(balloon->hostmem);
448 
449     our_range_memslots = hvb_our_range_memslots_new(balloon->addr,
450                                                     balloon->mr, hostmem_mr,
451                                                     OBJECT(balloon),
452                                                     balloon->memslot_count,
453                                                     balloon->memslot_size);
454     our_range = OUR_RANGE(our_range_memslots);
455 
456     if (hvb_page_range_tree_intree_any(balloon->removed_guest,
457                                        our_range->range.start,
458                                        our_range->range.count) ||
459         hvb_page_range_tree_intree_any(balloon->removed_both,
460                                        our_range->range.start,
461                                        our_range->range.count)) {
462         error_report("some parts of the memory backend were already returned by the guest. this should not happen, please reboot the guest and try again");
463         return false;
464     }
465 
466     trace_hv_balloon_our_range_add(our_range->range.count,
467                                    our_range->range.start);
468 
469     balloon->our_range = g_steal_pointer(&our_range_memslots);
470     return true;
471 }
472 
473 static void hv_balloon_hot_add_setup(HvBalloon *balloon, StateDesc *stdesc)
474 {
475     /* need to make copy since it is in union with hot_add_range */
476     uint64_t hot_add_diff = balloon->hot_add_diff;
477     PageRange *hot_add_range = &balloon->hot_add_range;
478     uint64_t align, our_range_remaining;
479     OurRange *our_range;
480 
481     assert(balloon->state == S_HOT_ADD_SETUP);
482     assert(hot_add_diff > 0);
483 
484     if (!hv_balloon_our_range_ensure(balloon)) {
485         goto ret_idle;
486     }
487 
488     our_range = OUR_RANGE(balloon->our_range);
489 
490     align = (1 << balloon->caps.cap_bits.hot_add_alignment) *
491         (MiB / HV_BALLOON_PAGE_SIZE);
492 
493     /* Absolute GPA in pages */
494     hot_add_range->start = our_range_get_remaining_start(our_range);
495     assert(QEMU_IS_ALIGNED(hot_add_range->start, align));
496 
497     our_range_remaining = our_range_get_remaining_size(our_range);
498     hot_add_range->count = MIN(our_range_remaining, hot_add_diff);
499     hot_add_range->count = QEMU_ALIGN_DOWN(hot_add_range->count, align);
500     if (hot_add_range->count == 0) {
501         goto ret_idle;
502     }
503 
504     hvb_our_range_memslots_ensure_mapped_additional(balloon->our_range,
505                                                     hot_add_range->count);
506 
507     HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_RB_WAIT);
508     return;
509 
510 ret_idle:
511     HV_BALLOON_STATE_DESC_SET(stdesc, S_IDLE);
512 }
513 
514 static void hv_balloon_hot_add_rb_wait(HvBalloon *balloon, StateDesc *stdesc)
515 {
516     VMBusChannel *chan = hv_balloon_get_channel(balloon);
517     struct dm_hot_add *ha;
518     size_t ha_size = sizeof(*ha) + sizeof(ha->range);
519 
520     assert(balloon->state == S_HOT_ADD_RB_WAIT);
521 
522     if (vmbus_channel_reserve(chan, 0, ha_size) < 0) {
523         return;
524     }
525 
526     HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_POSTING);
527 }
528 
529 static void hv_balloon_hot_add_posting(HvBalloon *balloon, StateDesc *stdesc)
530 {
531     PageRange *hot_add_range = &balloon->hot_add_range;
532     uint64_t *current_count = &balloon->ha_current_count;
533     VMBusChannel *chan = hv_balloon_get_channel(balloon);
534     struct dm_hot_add *ha;
535     size_t ha_size = sizeof(*ha) + sizeof(ha->range);
536     union dm_mem_page_range *ha_region;
537     uint64_t align, chunk_max_size;
538     ssize_t ret;
539 
540     assert(balloon->state == S_HOT_ADD_POSTING);
541     assert(hot_add_range->count > 0);
542 
543     align = (1 << balloon->caps.cap_bits.hot_add_alignment) *
544         (MiB / HV_BALLOON_PAGE_SIZE);
545     if (align >= HV_BALLOON_HA_CHUNK_PAGES) {
546         /*
547          * If the required alignment is higher than the chunk size we let it
548          * override that size.
549          */
550         chunk_max_size = align;
551     } else {
552         chunk_max_size = QEMU_ALIGN_DOWN(HV_BALLOON_HA_CHUNK_PAGES, align);
553     }
554 
555     /*
556      * hot_add_range->count starts aligned in hv_balloon_hot_add_setup(),
557      * then it is either reduced by subtracting aligned current_count or
558      * further hot-adds are prevented by marking the whole remaining our range
559      * as unusable in hv_balloon_handle_hot_add_response().
560      */
561     *current_count = MIN(hot_add_range->count, chunk_max_size);
562 
563     ha = alloca(ha_size);
564     ha_region = &(&ha->range)[1];
565     memset(ha, 0, ha_size);
566     ha->hdr.type = DM_MEM_HOT_ADD_REQUEST;
567     ha->hdr.size = ha_size;
568     ha->hdr.trans_id = balloon->trans_id;
569 
570     ha->range.finfo.start_page = hot_add_range->start;
571     ha->range.finfo.page_cnt = *current_count;
572     ha_region->finfo.start_page = hot_add_range->start;
573     ha_region->finfo.page_cnt = ha->range.finfo.page_cnt;
574 
575     trace_hv_balloon_outgoing_hot_add(ha->hdr.trans_id,
576                                       *current_count, hot_add_range->start);
577 
578     ret = vmbus_channel_send(chan, VMBUS_PACKET_DATA_INBAND,
579                              NULL, 0, ha, ha_size, false,
580                              ha->hdr.trans_id);
581     if (ret <= 0) {
582         error_report("error %zd when posting hot add msg, expect problems",
583                      ret);
584     }
585 
586     HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_REPLY_WAIT);
587 }
588 
589 static void hv_balloon_balloon_rb_wait(HvBalloon *balloon, StateDesc *stdesc)
590 {
591     VMBusChannel *chan = hv_balloon_get_channel(balloon);
592     size_t bl_size = sizeof(struct dm_balloon);
593 
594     assert(balloon->state == S_BALLOON_RB_WAIT);
595 
596     if (vmbus_channel_reserve(chan, 0, bl_size) < 0) {
597         return;
598     }
599 
600     HV_BALLOON_STATE_DESC_SET(stdesc, S_BALLOON_POSTING);
601 }
602 
603 static void hv_balloon_balloon_posting(HvBalloon *balloon, StateDesc *stdesc)
604 {
605     VMBusChannel *chan = hv_balloon_get_channel(balloon);
606     struct dm_balloon bl;
607     size_t bl_size = sizeof(bl);
608     ssize_t ret;
609 
610     assert(balloon->state == S_BALLOON_POSTING);
611     assert(balloon->balloon_diff > 0);
612 
613     memset(&bl, 0, sizeof(bl));
614     bl.hdr.type = DM_BALLOON_REQUEST;
615     bl.hdr.size = bl_size;
616     bl.hdr.trans_id = balloon->trans_id;
617     bl.num_pages = MIN(balloon->balloon_diff, HV_BALLOON_HR_CHUNK_PAGES);
618 
619     trace_hv_balloon_outgoing_balloon(bl.hdr.trans_id, bl.num_pages,
620                                       balloon->balloon_diff);
621 
622     ret = vmbus_channel_send(chan, VMBUS_PACKET_DATA_INBAND,
623                              NULL, 0, &bl, bl_size, false,
624                              bl.hdr.trans_id);
625     if (ret <= 0) {
626         error_report("error %zd when posting balloon msg, expect problems",
627                      ret);
628     }
629 
630     HV_BALLOON_STATE_DESC_SET(stdesc, S_BALLOON_REPLY_WAIT);
631 }
632 
633 static void hv_balloon_idle_state_process_target(HvBalloon *balloon,
634                                                  StateDesc *stdesc)
635 {
636     bool can_balloon = balloon->caps.cap_bits.balloon;
637     uint64_t ram_size_pages, total_removed;
638 
639     ram_size_pages = hv_balloon_total_ram(balloon);
640     total_removed = hv_balloon_total_removed_rs(balloon, ram_size_pages);
641 
642     /*
643      * we need to cache the values computed from the balloon target value when
644      * starting the adjustment procedure in case someone changes the target when
645      * the procedure is in progress
646      */
647     if (balloon->target > ram_size_pages - total_removed) {
648         bool can_hot_add = balloon->caps.cap_bits.hot_add;
649         uint64_t target_diff = balloon->target -
650             (ram_size_pages - total_removed);
651 
652         balloon->unballoon_diff = MIN(target_diff, total_removed);
653 
654         if (can_hot_add) {
655             balloon->hot_add_diff = target_diff - balloon->unballoon_diff;
656         } else {
657             balloon->hot_add_diff = 0;
658         }
659 
660         if (balloon->unballoon_diff > 0) {
661             assert(can_balloon);
662             HV_BALLOON_STATE_DESC_SET(stdesc, S_UNBALLOON_RB_WAIT);
663         } else if (balloon->hot_add_diff > 0) {
664             HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_SETUP);
665         }
666     } else if (can_balloon &&
667                balloon->target < ram_size_pages - total_removed) {
668         balloon->balloon_diff = ram_size_pages - total_removed -
669             balloon->target;
670         HV_BALLOON_STATE_DESC_SET(stdesc, S_BALLOON_RB_WAIT);
671     }
672 }
673 
674 static void hv_balloon_idle_state(HvBalloon *balloon,
675                                   StateDesc *stdesc)
676 {
677     assert(balloon->state == S_IDLE);
678 
679     if (balloon->target_changed) {
680         balloon->target_changed = false;
681         hv_balloon_idle_state_process_target(balloon, stdesc);
682         return;
683     }
684 }
685 
686 static const struct {
687     void (*handler)(HvBalloon *balloon, StateDesc *stdesc);
688 } state_handlers[] = {
689     [S_IDLE].handler = hv_balloon_idle_state,
690     [S_BALLOON_POSTING].handler = hv_balloon_balloon_posting,
691     [S_BALLOON_RB_WAIT].handler = hv_balloon_balloon_rb_wait,
692     [S_UNBALLOON_POSTING].handler = hv_balloon_unballoon_posting,
693     [S_UNBALLOON_RB_WAIT].handler = hv_balloon_unballoon_rb_wait,
694     [S_HOT_ADD_SETUP].handler = hv_balloon_hot_add_setup,
695     [S_HOT_ADD_RB_WAIT].handler = hv_balloon_hot_add_rb_wait,
696     [S_HOT_ADD_POSTING].handler = hv_balloon_hot_add_posting,
697 };
698 
699 static void hv_balloon_handle_state(HvBalloon *balloon, StateDesc *stdesc)
700 {
701     if (balloon->state >= ARRAY_SIZE(state_handlers) ||
702         !state_handlers[balloon->state].handler) {
703         return;
704     }
705 
706     state_handlers[balloon->state].handler(balloon, stdesc);
707 }
708 
709 static void hv_balloon_remove_response_insert_range(PageRangeTree tree,
710                                                     const PageRange *range,
711                                                     uint64_t *ctr1,
712                                                     uint64_t *ctr2,
713                                                     uint64_t *ctr3)
714 {
715     uint64_t dupcount, effcount;
716 
717     if (range->count == 0) {
718         return;
719     }
720 
721     dupcount = 0;
722     hvb_page_range_tree_insert(tree, range->start, range->count, &dupcount);
723 
724     assert(dupcount <= range->count);
725     effcount = range->count - dupcount;
726 
727     *ctr1 += effcount;
728     *ctr2 += effcount;
729     if (ctr3) {
730         *ctr3 += effcount;
731     }
732 }
733 
734 static void hv_balloon_remove_response_handle_range(HvBalloon *balloon,
735                                                     PageRange *range,
736                                                     bool both,
737                                                     uint64_t *removedctr)
738 {
739     OurRange *our_range = OUR_RANGE(balloon->our_range);
740     PageRangeTree globaltree =
741         both ? balloon->removed_both : balloon->removed_guest;
742     uint64_t *globalctr =
743         both ? &balloon->removed_both_ctr : &balloon->removed_guest_ctr;
744     PageRange rangeeff;
745 
746     if (range->count == 0) {
747         return;
748     }
749 
750     trace_hv_balloon_remove_response(range->count, range->start, both);
751 
752     if (our_range) {
753         /* Includes the not-yet-hot-added and unusable parts. */
754         rangeeff = our_range->range;
755     } else {
756         rangeeff.start = rangeeff.count = 0;
757     }
758 
759     if (page_range_intersection_size(range, rangeeff.start, rangeeff.count) > 0) {
760         PageRangeTree ourtree = our_range_get_removed_tree(our_range, both);
761         PageRange rangehole, rangecommon;
762         uint64_t ourremoved = 0;
763 
764         /* process the hole before our range, if it exists */
765         page_range_part_before(range, rangeeff.start, &rangehole);
766         hv_balloon_remove_response_insert_range(globaltree, &rangehole,
767                                                 globalctr, removedctr, NULL);
768         if (rangehole.count > 0) {
769             trace_hv_balloon_remove_response_hole(rangehole.count,
770                                                   rangehole.start,
771                                                   range->count, range->start,
772                                                   rangeeff.start, both);
773         }
774 
775         /* process our part */
776         page_range_intersect(range, rangeeff.start, rangeeff.count,
777                              &rangecommon);
778         hv_balloon_remove_response_insert_range(ourtree, &rangecommon,
779                                                 globalctr, removedctr,
780                                                 &ourremoved);
781         if (rangecommon.count > 0) {
782             trace_hv_balloon_remove_response_common(rangecommon.count,
783                                                     rangecommon.start,
784                                                     range->count, range->start,
785                                                     rangeeff.count,
786                                                     rangeeff.start, ourremoved,
787                                                     both);
788         }
789 
790         /* calculate what's left after our range */
791         rangecommon = *range;
792         page_range_part_after(&rangecommon, rangeeff.start, rangeeff.count,
793                               range);
794     }
795 
796     /* process the remainder of the range that lies after our range */
797     if (range->count > 0) {
798         hv_balloon_remove_response_insert_range(globaltree, range,
799                                                 globalctr, removedctr, NULL);
800         trace_hv_balloon_remove_response_remainder(range->count, range->start,
801                                                    both);
802         range->count = 0;
803     }
804 }
805 
806 static void hv_balloon_remove_response_handle_pages(HvBalloon *balloon,
807                                                     PageRange *range,
808                                                     uint64_t start,
809                                                     uint64_t count,
810                                                     bool both,
811                                                     uint64_t *removedctr)
812 {
813     assert(count > 0);
814 
815     /*
816      * if there is an existing range that the new range can't be joined to
817      * dump it into tree(s)
818      */
819     if (range->count > 0 && !page_range_joinable(range, start, count)) {
820         hv_balloon_remove_response_handle_range(balloon, range, both,
821                                                 removedctr);
822     }
823 
824     if (range->count == 0) {
825         range->start = start;
826         range->count = count;
827     } else if (page_range_joinable_left(range, start, count)) {
828         range->start = start;
829         range->count += count;
830     } else { /* page_range_joinable_right() */
831         range->count += count;
832     }
833 }
834 
835 static gboolean hv_balloon_handle_remove_host_addr_node(gpointer key,
836                                                         gpointer value,
837                                                         gpointer data)
838 {
839     PageRange *range = value;
840     uint64_t pageoff;
841 
842     for (pageoff = 0; pageoff < range->count; ) {
843         uint64_t addr_64 = (range->start + pageoff) * HV_BALLOON_PAGE_SIZE;
844         void *addr;
845         RAMBlock *rb;
846         ram_addr_t rb_offset;
847         size_t rb_page_size;
848         size_t discard_size;
849 
850         assert(addr_64 <= UINTPTR_MAX);
851         addr = (void *)((uintptr_t)addr_64);
852         rb = qemu_ram_block_from_host(addr, false, &rb_offset);
853         rb_page_size = qemu_ram_pagesize(rb);
854 
855         if (rb_page_size != HV_BALLOON_PAGE_SIZE) {
856             /* TODO: these should end in "removed_guest" */
857             warn_report("guest reported removed page backed by unsupported page size %zu",
858                         rb_page_size);
859             pageoff++;
860             continue;
861         }
862 
863         discard_size = MIN(range->count - pageoff,
864                            (rb->max_length - rb_offset) /
865                            HV_BALLOON_PAGE_SIZE);
866         discard_size = MAX(discard_size, 1);
867 
868         if (ram_block_discard_range(rb, rb_offset, discard_size *
869                                     HV_BALLOON_PAGE_SIZE) != 0) {
870             warn_report("guest reported removed page failed discard");
871         }
872 
873         pageoff += discard_size;
874     }
875 
876     return false;
877 }
878 
879 static void hv_balloon_handle_remove_host_addr_tree(PageRangeTree tree)
880 {
881     g_tree_foreach(tree.t, hv_balloon_handle_remove_host_addr_node, NULL);
882 }
883 
884 static int hv_balloon_handle_remove_section(PageRangeTree tree,
885                                             const MemoryRegionSection *section,
886                                             uint64_t count)
887 {
888     void *addr = memory_region_get_ram_ptr(section->mr) +
889         section->offset_within_region;
890     uint64_t addr_page;
891 
892     assert(count > 0);
893 
894     if ((uintptr_t)addr % HV_BALLOON_PAGE_SIZE) {
895         warn_report("guest reported removed pages at an unaligned host addr %p",
896                     addr);
897         return -EINVAL;
898     }
899 
900     addr_page = (uintptr_t)addr / HV_BALLOON_PAGE_SIZE;
901     hvb_page_range_tree_insert(tree, addr_page, count, NULL);
902 
903     return 0;
904 }
905 
906 static void hv_balloon_handle_remove_ranges(HvBalloon *balloon,
907                                             union dm_mem_page_range ranges[],
908                                             uint32_t count)
909 {
910     uint64_t removedcnt;
911     PageRangeTree removed_host_addr;
912     PageRange range_guest, range_both;
913 
914     hvb_page_range_tree_init(&removed_host_addr);
915     range_guest.count = range_both.count = removedcnt = 0;
916     for (unsigned int ctr = 0; ctr < count; ctr++) {
917         union dm_mem_page_range *mr = &ranges[ctr];
918         hwaddr pa;
919         MemoryRegionSection section;
920 
921         for (unsigned int offset = 0; offset < mr->finfo.page_cnt; ) {
922             int ret;
923             uint64_t pageno = mr->finfo.start_page + offset;
924             uint64_t pagecnt = 1;
925 
926             pa = (hwaddr)pageno << HV_BALLOON_PFN_SHIFT;
927             section = memory_region_find(get_system_memory(), pa,
928                                          (mr->finfo.page_cnt - offset) *
929                                          HV_BALLOON_PAGE_SIZE);
930             if (!section.mr) {
931                 warn_report("guest reported removed page %"PRIu64" not found in RAM",
932                             pageno);
933                 ret = -EINVAL;
934                 goto finish_page;
935             }
936 
937             pagecnt = int128_get64(section.size) / HV_BALLOON_PAGE_SIZE;
938             if (pagecnt <= 0) {
939                 warn_report("guest reported removed page %"PRIu64" in a section smaller than page size",
940                             pageno);
941                 pagecnt = 1; /* skip the whole page */
942                 ret = -EINVAL;
943                 goto finish_page;
944             }
945 
946             if (!memory_region_is_ram(section.mr) ||
947                 memory_region_is_rom(section.mr) ||
948                 memory_region_is_romd(section.mr)) {
949                 warn_report("guest reported removed page %"PRIu64" in a section that is not an ordinary RAM",
950                             pageno);
951                 ret = -EINVAL;
952                 goto finish_page;
953             }
954 
955             ret = hv_balloon_handle_remove_section(removed_host_addr, &section,
956                                                    pagecnt);
957 
958         finish_page:
959             if (ret == 0) {
960                 hv_balloon_remove_response_handle_pages(balloon,
961                                                         &range_both,
962                                                         pageno, pagecnt,
963                                                         true, &removedcnt);
964             } else {
965                 hv_balloon_remove_response_handle_pages(balloon,
966                                                         &range_guest,
967                                                         pageno, pagecnt,
968                                                         false, &removedcnt);
969             }
970 
971             if (section.mr) {
972                 memory_region_unref(section.mr);
973             }
974 
975             offset += pagecnt;
976         }
977     }
978 
979     hv_balloon_remove_response_handle_range(balloon, &range_both, true,
980                                             &removedcnt);
981     hv_balloon_remove_response_handle_range(balloon, &range_guest, false,
982                                             &removedcnt);
983 
984     hv_balloon_handle_remove_host_addr_tree(removed_host_addr);
985     hvb_page_range_tree_destroy(&removed_host_addr);
986 
987     if (removedcnt > balloon->balloon_diff) {
988         warn_report("guest reported more pages removed than currently pending (%"PRIu64" vs %"PRIu64")",
989                     removedcnt, balloon->balloon_diff);
990         balloon->balloon_diff = 0;
991     } else {
992         balloon->balloon_diff -= removedcnt;
993     }
994 }
995 
996 static bool hv_balloon_handle_msg_size(HvBalloonReq *req, size_t minsize,
997                                        const char *msgname)
998 {
999     VMBusChanReq *vmreq = &req->vmreq;
1000     uint32_t msglen = vmreq->msglen;
1001 
1002     if (msglen >= minsize) {
1003         return true;
1004     }
1005 
1006     warn_report("%s message too short (%u vs %zu), ignoring", msgname,
1007                 (unsigned int)msglen, minsize);
1008     return false;
1009 }
1010 
1011 static void hv_balloon_handle_version_request(HvBalloon *balloon,
1012                                               HvBalloonReq *req,
1013                                               StateDesc *stdesc)
1014 {
1015     VMBusChanReq *vmreq = &req->vmreq;
1016     struct dm_version_request *msgVr = vmreq->msg;
1017     struct dm_version_response respVr;
1018 
1019     if (balloon->state != S_VERSION) {
1020         warn_report("unexpected DM_VERSION_REQUEST in %d state",
1021                     balloon->state);
1022         return;
1023     }
1024 
1025     if (!hv_balloon_handle_msg_size(req, sizeof(*msgVr),
1026                                     "DM_VERSION_REQUEST")) {
1027         return;
1028     }
1029 
1030     trace_hv_balloon_incoming_version(msgVr->version.major_version,
1031                                       msgVr->version.minor_version);
1032 
1033     memset(&respVr, 0, sizeof(respVr));
1034     respVr.hdr.type = DM_VERSION_RESPONSE;
1035     respVr.hdr.size = sizeof(respVr);
1036     respVr.hdr.trans_id = msgVr->hdr.trans_id;
1037     respVr.is_accepted = msgVr->version.version >= DYNMEM_PROTOCOL_VERSION_1 &&
1038         msgVr->version.version <= DYNMEM_PROTOCOL_VERSION_3;
1039 
1040     hv_balloon_send_packet(vmreq->chan, (struct dm_message *)&respVr);
1041 
1042     if (respVr.is_accepted) {
1043         HV_BALLOON_STATE_DESC_SET(stdesc, S_CAPS);
1044     }
1045 }
1046 
1047 static void hv_balloon_handle_caps_report(HvBalloon *balloon,
1048                                           HvBalloonReq *req,
1049                                           StateDesc *stdesc)
1050 {
1051     VMBusChanReq *vmreq = &req->vmreq;
1052     struct dm_capabilities *msgCap = vmreq->msg;
1053     struct dm_capabilities_resp_msg respCap;
1054 
1055     if (balloon->state != S_CAPS) {
1056         warn_report("unexpected DM_CAPABILITIES_REPORT in %d state",
1057                     balloon->state);
1058         return;
1059     }
1060 
1061     if (!hv_balloon_handle_msg_size(req, sizeof(*msgCap),
1062                                     "DM_CAPABILITIES_REPORT")) {
1063         return;
1064     }
1065 
1066     trace_hv_balloon_incoming_caps(msgCap->caps.caps);
1067     balloon->caps = msgCap->caps;
1068 
1069     memset(&respCap, 0, sizeof(respCap));
1070     respCap.hdr.type = DM_CAPABILITIES_RESPONSE;
1071     respCap.hdr.size = sizeof(respCap);
1072     respCap.hdr.trans_id = msgCap->hdr.trans_id;
1073     respCap.is_accepted = 1;
1074     respCap.hot_remove = 1;
1075     respCap.suppress_pressure_reports = !balloon->status_report.enabled;
1076     hv_balloon_send_packet(vmreq->chan, (struct dm_message *)&respCap);
1077 
1078     timer_mod(&balloon->post_init_timer,
1079               qemu_clock_get_ms(QEMU_CLOCK_VIRTUAL) +
1080               HV_BALLOON_POST_INIT_WAIT);
1081 
1082     HV_BALLOON_STATE_DESC_SET(stdesc, S_POST_INIT_WAIT);
1083 }
1084 
1085 static void hv_balloon_handle_status_report(HvBalloon *balloon,
1086                                             HvBalloonReq *req)
1087 {
1088     VMBusChanReq *vmreq = &req->vmreq;
1089     struct dm_status *msgStatus = vmreq->msg;
1090 
1091     if (!hv_balloon_handle_msg_size(req, sizeof(*msgStatus),
1092                                     "DM_STATUS_REPORT")) {
1093         return;
1094     }
1095 
1096     if (!balloon->status_report.enabled) {
1097         return;
1098     }
1099 
1100     balloon->status_report.committed = msgStatus->num_committed;
1101     balloon->status_report.committed *= HV_BALLOON_PAGE_SIZE;
1102     balloon->status_report.available = msgStatus->num_avail;
1103     balloon->status_report.available *= HV_BALLOON_PAGE_SIZE;
1104     balloon->status_report.received = true;
1105 
1106     qapi_event_send_hv_balloon_status_report(balloon->status_report.committed,
1107                                              balloon->status_report.available);
1108 }
1109 
1110 HvBalloonInfo *qmp_query_hv_balloon_status_report(Error **errp)
1111 {
1112     HvBalloon *balloon;
1113     HvBalloonInfo *info;
1114 
1115     balloon = HV_BALLOON(object_resolve_path_type("", TYPE_HV_BALLOON, NULL));
1116     if (!balloon) {
1117         error_setg(errp, "no %s device present", TYPE_HV_BALLOON);
1118         return NULL;
1119     }
1120 
1121     if (!balloon->status_report.enabled) {
1122         error_setg(errp, "guest memory status reporting not enabled");
1123         return NULL;
1124     }
1125 
1126     if (!balloon->status_report.received) {
1127         error_setg(errp, "no guest memory status report received yet");
1128         return NULL;
1129     }
1130 
1131     info = g_malloc0(sizeof(*info));
1132     info->committed = balloon->status_report.committed;
1133     info->available = balloon->status_report.available;
1134     return info;
1135 }
1136 
1137 static void hv_balloon_handle_unballoon_response(HvBalloon *balloon,
1138                                                  HvBalloonReq *req,
1139                                                  StateDesc *stdesc)
1140 {
1141     VMBusChanReq *vmreq = &req->vmreq;
1142     struct dm_unballoon_response *msgUrR = vmreq->msg;
1143 
1144     if (balloon->state != S_UNBALLOON_REPLY_WAIT) {
1145         warn_report("unexpected DM_UNBALLOON_RESPONSE in %d state",
1146                     balloon->state);
1147         return;
1148     }
1149 
1150     if (!hv_balloon_handle_msg_size(req, sizeof(*msgUrR),
1151                                     "DM_UNBALLOON_RESPONSE"))
1152         return;
1153 
1154     trace_hv_balloon_incoming_unballoon(msgUrR->hdr.trans_id);
1155 
1156     balloon->trans_id++;
1157 
1158     if (balloon->hot_add_diff > 0) {
1159         bool can_hot_add = balloon->caps.cap_bits.hot_add;
1160 
1161         assert(can_hot_add);
1162         HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_SETUP);
1163     } else {
1164         HV_BALLOON_STATE_DESC_SET(stdesc, S_IDLE);
1165     }
1166 }
1167 
1168 static void hv_balloon_handle_hot_add_response(HvBalloon *balloon,
1169                                                HvBalloonReq *req,
1170                                                StateDesc *stdesc)
1171 {
1172     PageRange *hot_add_range = &balloon->hot_add_range;
1173     VMBusChanReq *vmreq = &req->vmreq;
1174     struct dm_hot_add_response *msgHaR = vmreq->msg;
1175     OurRange *our_range;
1176 
1177     if (balloon->state != S_HOT_ADD_REPLY_WAIT) {
1178         warn_report("unexpected DM_HOT_ADD_RESPONSE in %d state",
1179                     balloon->state);
1180         return;
1181     }
1182 
1183     assert(balloon->our_range);
1184     our_range = OUR_RANGE(balloon->our_range);
1185 
1186     if (!hv_balloon_handle_msg_size(req, sizeof(*msgHaR),
1187                                     "DM_HOT_ADD_RESPONSE"))
1188         return;
1189 
1190     trace_hv_balloon_incoming_hot_add(msgHaR->hdr.trans_id, msgHaR->result,
1191                                       msgHaR->page_count);
1192 
1193     balloon->trans_id++;
1194 
1195     if (msgHaR->result) {
1196         if (msgHaR->page_count > balloon->ha_current_count) {
1197             warn_report("DM_HOT_ADD_RESPONSE page count higher than requested (%"PRIu32" vs %"PRIu64")",
1198                         msgHaR->page_count, balloon->ha_current_count);
1199             msgHaR->page_count = balloon->ha_current_count;
1200         }
1201 
1202         hvb_our_range_mark_added(our_range, msgHaR->page_count);
1203         hot_add_range->start += msgHaR->page_count;
1204         hot_add_range->count -= msgHaR->page_count;
1205     }
1206 
1207     if (!msgHaR->result || msgHaR->page_count < balloon->ha_current_count) {
1208         /*
1209          * the current planned range was only partially hot-added, take note
1210          * how much of it remains and don't attempt any further hot adds
1211          */
1212         our_range_mark_remaining_unusable(our_range);
1213 
1214         goto ret_idle;
1215     }
1216 
1217     /* any pages remaining to hot-add in our range? */
1218     if (hot_add_range->count > 0) {
1219         HV_BALLOON_STATE_DESC_SET(stdesc, S_HOT_ADD_RB_WAIT);
1220         return;
1221     }
1222 
1223 ret_idle:
1224     HV_BALLOON_STATE_DESC_SET(stdesc, S_IDLE);
1225 }
1226 
1227 static void hv_balloon_handle_balloon_response(HvBalloon *balloon,
1228                                                HvBalloonReq *req,
1229                                                StateDesc *stdesc)
1230 {
1231     VMBusChanReq *vmreq = &req->vmreq;
1232     struct dm_balloon_response *msgBR = vmreq->msg;
1233 
1234     if (balloon->state != S_BALLOON_REPLY_WAIT) {
1235         warn_report("unexpected DM_BALLOON_RESPONSE in %d state",
1236                     balloon->state);
1237         return;
1238     }
1239 
1240     if (!hv_balloon_handle_msg_size(req, sizeof(*msgBR),
1241                                     "DM_BALLOON_RESPONSE"))
1242         return;
1243 
1244     trace_hv_balloon_incoming_balloon(msgBR->hdr.trans_id, msgBR->range_count,
1245                                       msgBR->more_pages);
1246 
1247     if (vmreq->msglen < sizeof(*msgBR) +
1248         (uint64_t)sizeof(msgBR->range_array[0]) * msgBR->range_count) {
1249         warn_report("DM_BALLOON_RESPONSE too short for the range count");
1250         return;
1251     }
1252 
1253     if (msgBR->range_count == 0) {
1254         /* The guest is already at its minimum size */
1255         balloon->balloon_diff = 0;
1256         goto ret_end_trans;
1257     } else {
1258         hv_balloon_handle_remove_ranges(balloon,
1259                                         msgBR->range_array,
1260                                         msgBR->range_count);
1261     }
1262 
1263     /* More responses expected? */
1264     if (msgBR->more_pages) {
1265         return;
1266     }
1267 
1268 ret_end_trans:
1269     balloon->trans_id++;
1270 
1271     if (balloon->balloon_diff > 0) {
1272         HV_BALLOON_STATE_DESC_SET(stdesc, S_BALLOON_RB_WAIT);
1273     } else {
1274         HV_BALLOON_STATE_DESC_SET(stdesc, S_IDLE);
1275     }
1276 }
1277 
1278 static void hv_balloon_handle_packet(HvBalloon *balloon, HvBalloonReq *req,
1279                                      StateDesc *stdesc)
1280 {
1281     VMBusChanReq *vmreq = &req->vmreq;
1282     struct dm_message *msg = vmreq->msg;
1283 
1284     if (vmreq->msglen < sizeof(msg->hdr)) {
1285         return;
1286     }
1287 
1288     switch (msg->hdr.type) {
1289     case DM_VERSION_REQUEST:
1290         hv_balloon_handle_version_request(balloon, req, stdesc);
1291         break;
1292 
1293     case DM_CAPABILITIES_REPORT:
1294         hv_balloon_handle_caps_report(balloon, req, stdesc);
1295         break;
1296 
1297     case DM_STATUS_REPORT:
1298         hv_balloon_handle_status_report(balloon, req);
1299         break;
1300 
1301     case DM_MEM_HOT_ADD_RESPONSE:
1302         hv_balloon_handle_hot_add_response(balloon, req, stdesc);
1303         break;
1304 
1305     case DM_UNBALLOON_RESPONSE:
1306         hv_balloon_handle_unballoon_response(balloon, req, stdesc);
1307         break;
1308 
1309     case DM_BALLOON_RESPONSE:
1310         hv_balloon_handle_balloon_response(balloon, req, stdesc);
1311         break;
1312 
1313     default:
1314         warn_report("unknown DM message %u", msg->hdr.type);
1315         break;
1316     }
1317 }
1318 
1319 static bool hv_balloon_recv_channel(HvBalloon *balloon, StateDesc *stdesc)
1320 {
1321     VMBusChannel *chan;
1322     HvBalloonReq *req;
1323 
1324     if (balloon->state == S_WAIT_RESET ||
1325         balloon->state == S_POST_RESET_CLOSED) {
1326         return false;
1327     }
1328 
1329     chan = hv_balloon_get_channel(balloon);
1330     if (vmbus_channel_recv_start(chan)) {
1331         return false;
1332     }
1333 
1334     while ((req = vmbus_channel_recv_peek(chan, sizeof(*req)))) {
1335         hv_balloon_handle_packet(balloon, req, stdesc);
1336         vmbus_free_req(req);
1337         vmbus_channel_recv_pop(chan);
1338 
1339         if (stdesc->state != S_NO_CHANGE) {
1340             break;
1341         }
1342     }
1343 
1344     return vmbus_channel_recv_done(chan) > 0;
1345 }
1346 
1347 /* old state handler -> new state transition (potential) */
1348 static bool hv_balloon_event_loop_state(HvBalloon *balloon)
1349 {
1350     StateDesc state_new = HV_BALLOON_STATE_DESC_INIT;
1351 
1352     hv_balloon_handle_state(balloon, &state_new);
1353     return hv_balloon_state_set(balloon, state_new.state, state_new.desc);
1354 }
1355 
1356 /* VMBus message -> new state transition (potential) */
1357 static bool hv_balloon_event_loop_recv(HvBalloon *balloon)
1358 {
1359     StateDesc state_new = HV_BALLOON_STATE_DESC_INIT;
1360     bool any_recv, state_changed;
1361 
1362     any_recv = hv_balloon_recv_channel(balloon, &state_new);
1363     state_changed = hv_balloon_state_set(balloon,
1364                                          state_new.state, state_new.desc);
1365 
1366     return state_changed || any_recv;
1367 }
1368 
1369 static void hv_balloon_event_loop(HvBalloon *balloon)
1370 {
1371     bool state_repeat, recv_repeat;
1372 
1373     do {
1374         state_repeat = hv_balloon_event_loop_state(balloon);
1375         recv_repeat = hv_balloon_event_loop_recv(balloon);
1376     } while (state_repeat || recv_repeat);
1377 }
1378 
1379 static void hv_balloon_vmdev_chan_notify(VMBusChannel *chan)
1380 {
1381     HvBalloon *balloon = HV_BALLOON(vmbus_channel_device(chan));
1382 
1383     hv_balloon_event_loop(balloon);
1384 }
1385 
1386 static void hv_balloon_stat(void *opaque, BalloonInfo *info)
1387 {
1388     HvBalloon *balloon = opaque;
1389     info->actual = (hv_balloon_total_ram(balloon) - balloon->removed_both_ctr)
1390         << HV_BALLOON_PFN_SHIFT;
1391 }
1392 
1393 static void hv_balloon_to_target(void *opaque, ram_addr_t target)
1394 {
1395     HvBalloon *balloon = opaque;
1396     uint64_t target_pages = target >> HV_BALLOON_PFN_SHIFT;
1397 
1398     if (!target_pages) {
1399         return;
1400     }
1401 
1402     /*
1403      * always set target_changed, even with unchanged target, as the user
1404      * might be asking us to try again reaching it
1405      */
1406     balloon->target = target_pages;
1407     balloon->target_changed = true;
1408 
1409     hv_balloon_event_loop(balloon);
1410 }
1411 
1412 static int hv_balloon_vmdev_open_channel(VMBusChannel *chan)
1413 {
1414     HvBalloon *balloon = HV_BALLOON(vmbus_channel_device(chan));
1415 
1416     if (balloon->state != S_POST_RESET_CLOSED) {
1417         warn_report("guest trying to open a DM channel in invalid %d state",
1418                     balloon->state);
1419         return -EINVAL;
1420     }
1421 
1422     HV_BALLOON_SET_STATE(balloon, S_VERSION);
1423     hv_balloon_event_loop(balloon);
1424 
1425     return 0;
1426 }
1427 
1428 static void hv_balloon_vmdev_close_channel(VMBusChannel *chan)
1429 {
1430     HvBalloon *balloon = HV_BALLOON(vmbus_channel_device(chan));
1431 
1432     timer_del(&balloon->post_init_timer);
1433 
1434     /* Don't report stale data */
1435     balloon->status_report.received = false;
1436 
1437     HV_BALLOON_SET_STATE(balloon, S_WAIT_RESET);
1438     hv_balloon_event_loop(balloon);
1439 }
1440 
1441 static void hv_balloon_post_init_timer(void *opaque)
1442 {
1443     HvBalloon *balloon = opaque;
1444 
1445     if (balloon->state != S_POST_INIT_WAIT) {
1446         return;
1447     }
1448 
1449     HV_BALLOON_SET_STATE(balloon, S_IDLE);
1450     hv_balloon_event_loop(balloon);
1451 }
1452 
1453 static void hv_balloon_system_reset_unrealize_common(HvBalloon *balloon)
1454 {
1455     g_clear_pointer(&balloon->our_range, hvb_our_range_memslots_free);
1456 }
1457 
1458 static void hv_balloon_system_reset(void *opaque)
1459 {
1460     HvBalloon *balloon = HV_BALLOON(opaque);
1461 
1462     hv_balloon_system_reset_unrealize_common(balloon);
1463 }
1464 
1465 static void hv_balloon_ensure_mr(HvBalloon *balloon)
1466 {
1467     MemoryRegion *hostmem_mr;
1468 
1469     assert(balloon->hostmem);
1470 
1471     if (balloon->mr) {
1472         return;
1473     }
1474 
1475     hostmem_mr = host_memory_backend_get_memory(balloon->hostmem);
1476 
1477     balloon->mr = g_new0(MemoryRegion, 1);
1478     memory_region_init(balloon->mr, OBJECT(balloon), TYPE_HV_BALLOON,
1479                        memory_region_size(hostmem_mr));
1480     balloon->mr->align = memory_region_get_alignment(hostmem_mr);
1481 }
1482 
1483 static void hv_balloon_free_mr(HvBalloon *balloon)
1484 {
1485     if (!balloon->mr) {
1486         return;
1487     }
1488 
1489     object_unparent(OBJECT(balloon->mr));
1490     g_clear_pointer(&balloon->mr, g_free);
1491 }
1492 
1493 static void hv_balloon_vmdev_realize(VMBusDevice *vdev, Error **errp)
1494 {
1495     ERRP_GUARD();
1496     HvBalloon *balloon = HV_BALLOON(vdev);
1497     int ret;
1498 
1499     balloon->state = S_WAIT_RESET;
1500 
1501     ret = qemu_add_balloon_handler(hv_balloon_to_target, hv_balloon_stat,
1502                                    balloon);
1503     if (ret < 0) {
1504         /* This also protects against having multiple hv-balloon instances */
1505         error_setg(errp, "Only one balloon device is supported");
1506         return;
1507     }
1508 
1509     if (balloon->hostmem) {
1510         if (host_memory_backend_is_mapped(balloon->hostmem)) {
1511             Object *obj = OBJECT(balloon->hostmem);
1512 
1513             error_setg(errp, "'%s' property specifies a busy memdev: %s",
1514                        HV_BALLOON_MEMDEV_PROP,
1515                        object_get_canonical_path_component(obj));
1516             goto out_balloon_handler;
1517         }
1518 
1519         hv_balloon_ensure_mr(balloon);
1520 
1521         /* This is rather unlikely to happen, but let's still check for it. */
1522         if (!QEMU_IS_ALIGNED(memory_region_size(balloon->mr),
1523                              HV_BALLOON_PAGE_SIZE)) {
1524             error_setg(errp, "'%s' property memdev size has to be a multiple of 0x%" PRIx64,
1525                        HV_BALLOON_MEMDEV_PROP, (uint64_t)HV_BALLOON_PAGE_SIZE);
1526             goto out_balloon_handler;
1527         }
1528 
1529         host_memory_backend_set_mapped(balloon->hostmem, true);
1530         vmstate_register_ram(host_memory_backend_get_memory(balloon->hostmem),
1531                              DEVICE(balloon));
1532     } else if (balloon->addr) {
1533         error_setg(errp, "'%s' property must not be set without a memdev",
1534                    HV_BALLOON_MEMDEV_PROP);
1535         goto out_balloon_handler;
1536     }
1537 
1538     timer_init_ms(&balloon->post_init_timer, QEMU_CLOCK_VIRTUAL,
1539                   hv_balloon_post_init_timer, balloon);
1540 
1541     qemu_register_reset(hv_balloon_system_reset, balloon);
1542 
1543     return;
1544 
1545 out_balloon_handler:
1546     qemu_remove_balloon_handler(balloon);
1547 }
1548 
1549 /*
1550  * VMBus device reset has to be implemented in case the guest decides to
1551  * disconnect and reconnect to the VMBus without rebooting the whole system.
1552  *
1553  * However, the hot-added memory can't be removed here as Windows keeps on using
1554  * it until the system is restarted, even after disconnecting from the VMBus.
1555  */
1556 static void hv_balloon_vmdev_reset(VMBusDevice *vdev)
1557 {
1558     HvBalloon *balloon = HV_BALLOON(vdev);
1559 
1560     if (balloon->state == S_POST_RESET_CLOSED) {
1561         return;
1562     }
1563 
1564     if (balloon->our_range) {
1565         hvb_our_range_clear_removed_trees(OUR_RANGE(balloon->our_range));
1566     }
1567 
1568     hvb_page_range_tree_destroy(&balloon->removed_guest);
1569     hvb_page_range_tree_destroy(&balloon->removed_both);
1570     hvb_page_range_tree_init(&balloon->removed_guest);
1571     hvb_page_range_tree_init(&balloon->removed_both);
1572 
1573     balloon->trans_id = 0;
1574     balloon->removed_guest_ctr = 0;
1575     balloon->removed_both_ctr = 0;
1576 
1577     HV_BALLOON_SET_STATE(balloon, S_POST_RESET_CLOSED);
1578     hv_balloon_event_loop(balloon);
1579 }
1580 
1581 /*
1582  * Clean up things that were (possibly) allocated pre-realization, for example
1583  * from memory_device_pre_plug(), so we don't leak them if the device don't
1584  * actually get realized in the end.
1585  */
1586 static void hv_balloon_unrealize_finalize_common(HvBalloon *balloon)
1587 {
1588     hv_balloon_free_mr(balloon);
1589     balloon->addr = 0;
1590 
1591     balloon->memslot_count = 0;
1592 }
1593 
1594 static void hv_balloon_vmdev_unrealize(VMBusDevice *vdev)
1595 {
1596     HvBalloon *balloon = HV_BALLOON(vdev);
1597 
1598     qemu_unregister_reset(hv_balloon_system_reset, balloon);
1599 
1600     hv_balloon_system_reset_unrealize_common(balloon);
1601 
1602     qemu_remove_balloon_handler(balloon);
1603 
1604     if (balloon->hostmem) {
1605         vmstate_unregister_ram(host_memory_backend_get_memory(balloon->hostmem),
1606                                DEVICE(balloon));
1607         host_memory_backend_set_mapped(balloon->hostmem, false);
1608     }
1609 
1610     hvb_page_range_tree_destroy(&balloon->removed_guest);
1611     hvb_page_range_tree_destroy(&balloon->removed_both);
1612 
1613     hv_balloon_unrealize_finalize_common(balloon);
1614 }
1615 
1616 static uint64_t hv_balloon_md_get_addr(const MemoryDeviceState *md)
1617 {
1618     return object_property_get_uint(OBJECT(md), HV_BALLOON_ADDR_PROP,
1619                                     &error_abort);
1620 }
1621 
1622 static void hv_balloon_md_set_addr(MemoryDeviceState *md, uint64_t addr,
1623                                    Error **errp)
1624 {
1625     object_property_set_uint(OBJECT(md), HV_BALLOON_ADDR_PROP, addr, errp);
1626 }
1627 
1628 static MemoryRegion *hv_balloon_md_get_memory_region(MemoryDeviceState *md,
1629                                                      Error **errp)
1630 {
1631     HvBalloon *balloon = HV_BALLOON(md);
1632 
1633     if (!balloon->hostmem) {
1634         return NULL;
1635     }
1636 
1637     hv_balloon_ensure_mr(balloon);
1638 
1639     return balloon->mr;
1640 }
1641 
1642 static uint64_t hv_balloon_md_get_min_alignment(const MemoryDeviceState *md)
1643 {
1644     /*
1645      * The VM can indicate an alignment up to 32 GiB. Memory device core can
1646      * usually only handle/guarantee 1 GiB alignment. The user will have to
1647      * specify a larger maxmem eventually.
1648      *
1649      * The memory device core will warn the user in case maxmem might have to be
1650      * increased and will fail plugging the device if there is not sufficient
1651      * space after alignment.
1652      *
1653      * TODO: we could do the alignment ourselves in a slightly bigger region.
1654      * But this feels better, although the warning might be annoying. Maybe
1655      * we can optimize that in the future (e.g., with such a device on the
1656      * cmdline place/size the device memory region differently.
1657      */
1658     return 32 * GiB;
1659 }
1660 
1661 static void hv_balloon_md_fill_device_info(const MemoryDeviceState *md,
1662                                            MemoryDeviceInfo *info)
1663 {
1664     HvBalloonDeviceInfo *hi = g_new0(HvBalloonDeviceInfo, 1);
1665     const HvBalloon *balloon = HV_BALLOON(md);
1666     DeviceState *dev = DEVICE(md);
1667 
1668     if (dev->id) {
1669         hi->id = g_strdup(dev->id);
1670     }
1671 
1672     if (balloon->hostmem) {
1673         hi->memdev = object_get_canonical_path(OBJECT(balloon->hostmem));
1674         hi->memaddr = balloon->addr;
1675         hi->has_memaddr = true;
1676         hi->max_size = memory_region_size(balloon->mr);
1677         /* TODO: expose current provided size or something else? */
1678     } else {
1679         hi->max_size = 0;
1680     }
1681 
1682     info->u.hv_balloon.data = hi;
1683     info->type = MEMORY_DEVICE_INFO_KIND_HV_BALLOON;
1684 }
1685 
1686 static void hv_balloon_decide_memslots(MemoryDeviceState *md,
1687                                        unsigned int limit)
1688 {
1689     HvBalloon *balloon = HV_BALLOON(md);
1690     MemoryRegion *hostmem_mr;
1691     uint64_t region_size, memslot_size, memslots;
1692 
1693     /* We're called exactly once, before realizing the device. */
1694     assert(!balloon->memslot_count);
1695 
1696     /* We should not be called if we don't have a memory backend */
1697     assert(balloon->hostmem);
1698 
1699     hostmem_mr = host_memory_backend_get_memory(balloon->hostmem);
1700     region_size = memory_region_size(hostmem_mr);
1701 
1702     assert(region_size > 0);
1703     memslot_size = QEMU_ALIGN_UP(region_size / limit,
1704                                  HV_BALLOON_HA_MEMSLOT_SIZE_ALIGN);
1705     memslots = QEMU_ALIGN_UP(region_size, memslot_size) / memslot_size;
1706 
1707     if (memslots > 1) {
1708         balloon->memslot_size = memslot_size;
1709     } else {
1710         balloon->memslot_size = region_size;
1711     }
1712 
1713     assert(memslots <= UINT_MAX);
1714     balloon->memslot_count = memslots;
1715 }
1716 
1717 static unsigned int hv_balloon_get_memslots(MemoryDeviceState *md)
1718 {
1719     const HvBalloon *balloon = HV_BALLOON(md);
1720 
1721     /* We're called after setting the suggested limit. */
1722     assert(balloon->memslot_count > 0);
1723 
1724     return balloon->memslot_count;
1725 }
1726 
1727 static void hv_balloon_init(Object *obj)
1728 {
1729 }
1730 
1731 static void hv_balloon_finalize(Object *obj)
1732 {
1733     HvBalloon *balloon = HV_BALLOON(obj);
1734 
1735     hv_balloon_unrealize_finalize_common(balloon);
1736 }
1737 
1738 static Property hv_balloon_properties[] = {
1739     DEFINE_PROP_BOOL("status-report", HvBalloon,
1740                      status_report.enabled, false),
1741 
1742     /* MEMORY_DEVICE props */
1743     DEFINE_PROP_LINK(HV_BALLOON_MEMDEV_PROP, HvBalloon, hostmem,
1744                      TYPE_MEMORY_BACKEND, HostMemoryBackend *),
1745     DEFINE_PROP_UINT64(HV_BALLOON_ADDR_PROP, HvBalloon, addr, 0),
1746 
1747     DEFINE_PROP_END_OF_LIST(),
1748 };
1749 
1750 static void hv_balloon_class_init(ObjectClass *klass, void *data)
1751 {
1752     DeviceClass *dc = DEVICE_CLASS(klass);
1753     VMBusDeviceClass *vdc = VMBUS_DEVICE_CLASS(klass);
1754     MemoryDeviceClass *mdc = MEMORY_DEVICE_CLASS(klass);
1755 
1756     device_class_set_props(dc, hv_balloon_properties);
1757     qemu_uuid_parse(HV_BALLOON_GUID, &vdc->classid);
1758     set_bit(DEVICE_CATEGORY_MISC, dc->categories);
1759 
1760     vdc->vmdev_realize = hv_balloon_vmdev_realize;
1761     vdc->vmdev_unrealize = hv_balloon_vmdev_unrealize;
1762     vdc->vmdev_reset = hv_balloon_vmdev_reset;
1763     vdc->open_channel = hv_balloon_vmdev_open_channel;
1764     vdc->close_channel = hv_balloon_vmdev_close_channel;
1765     vdc->chan_notify_cb = hv_balloon_vmdev_chan_notify;
1766 
1767     mdc->get_addr = hv_balloon_md_get_addr;
1768     mdc->set_addr = hv_balloon_md_set_addr;
1769     mdc->get_plugged_size = memory_device_get_region_size;
1770     mdc->get_memory_region = hv_balloon_md_get_memory_region;
1771     mdc->decide_memslots = hv_balloon_decide_memslots;
1772     mdc->get_memslots = hv_balloon_get_memslots;
1773     mdc->get_min_alignment = hv_balloon_md_get_min_alignment;
1774     mdc->fill_device_info = hv_balloon_md_fill_device_info;
1775 }
1776