Copy in style scheduler changes

Author: kseamon

src/cdk/table/coalesced-style-scheduler.ts

@@ -6,8 +6,17 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {Injectable, NgZone} from '@angular/core';
-import {from, Observable} from 'rxjs';
+import {Injectable, NgZone, OnDestroy} from '@angular/core';
+import {Subject} from 'rxjs';
+import {take, takeUntil} from 'rxjs/operators';
+
+/**
+ * @docs-private
+ */
+export class _Schedule {
+  tasks: (() => unknown)[] = [];
+  endTasks: (() => unknown)[] = [];
+}
 
 /**
  * Allows grouping up CSSDom mutations after the current execution context.
@@ -17,40 +26,55 @@ import {from, Observable} from 'rxjs';
  * @docs-private
  */
 @Injectable()
-export class _CoalescedStyleScheduler {
-  private _currentSchedule: Observable<void>|null = null;
+export class _CoalescedStyleScheduler implements OnDestroy {
+  private _currentSchedule: _Schedule|null = null;
+  private readonly _destroyed = new Subject<void>();
 
   constructor(private readonly _ngZone: NgZone) {}
 
   /**
-   * Schedules the specified task to run after the current microtask.
+   * Schedules the specified task to run at the end of the current VM turn.
    */
   schedule(task: () => unknown): void {
     this._createScheduleIfNeeded();
 
-    this._currentSchedule!.subscribe(task);
+    this._currentSchedule!.tasks.push(task);
   }
 
   /**
-   * Schedules the specified task to run after all scheduled tasks.
-   * This is useful when measuring after style updates is required.
+   * Schedules the specified task to run after other scheduled tasks at the end of the current
+   * VM turn.
    */
   scheduleEnd(task: () => unknown): void {
     this._createScheduleIfNeeded();
 
-    this._currentSchedule!.subscribe(() => {
-      this.schedule(task);
-    });
+    this._currentSchedule!.endTasks.push(task);
+  }
+
+  /** Prevent any further tasks from running. */
+  ngOnDestroy() {
+    this._destroyed.next();
+    this._destroyed.complete();
   }
 
   private _createScheduleIfNeeded() {
     if (this._currentSchedule) { return; }
 
-    this._ngZone.runOutsideAngular(() => {
-      this._currentSchedule = from(new Promise<void>((resolve) => {
-        this._currentSchedule = null;
-        resolve(undefined);
-      }));
+    this._currentSchedule = new _Schedule();
+
+    this._ngZone.onStable.pipe(
+        take(1),
+        takeUntil(this._destroyed),
+    ).subscribe(() => {
+      const schedule = this._currentSchedule!;
+      this._currentSchedule = null;
+
+      for (const task of schedule.tasks) {
+        task();
+      }
+      for (const task of schedule.endTasks) {
+        task();
+      }
     });
   }
 }