[docs] [issue31139] Improve tracemalloc demo code in docs