diff --git a/polygerrit-ui/app/api/annotation.ts b/polygerrit-ui/app/api/annotation.ts
index c046b4f..e58bdd5 100644
--- a/polygerrit-ui/app/api/annotation.ts
+++ b/polygerrit-ui/app/api/annotation.ts
@@ -14,18 +14,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import {CoverageRange, Side} from './diff';
+import {CoverageRange, GrDiffLine, Side} from './diff';
 import {StyleObject} from './styles';
 
-export type AddLayerFunc = (ctx: AnnotationContext) => void;
-
-export type NotifyFunc = (
-  path: string,
-  start: number,
-  end: number,
-  side: Side
-) => void;
-
+/**
+ * This is the callback object that Gerrit calls once for each diff. Gerrit
+ * is then responsible for styling the diff according the returned array of
+ * CoverageRanges.
+ */
 export type CoverageProvider = (
   changeNum: number,
   path: string,
@@ -34,14 +30,35 @@
   /**
    * This is a ChangeInfo object as defined here:
    * https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#change-info
-   * We neither want to repeat it nor add a dependency on it here.
+   * At the moment we neither want to repeat it nor add a dependency on it here.
+   * TODO: Create a dedicated smaller object for exposing a change in the plugin
+   * API. Or allow the plugin API to depend on the entire rest API.
    */
   change?: unknown
 ) => Promise<Array<CoverageRange>>;
 
+export type AnnotationCallback = (ctx: AnnotationContext) => void;
+
+/**
+ * This object is passed to the plugin from Gerrit for each line of a diff that
+ * is being rendered. The plugin can then call annotateRange() or
+ * annotateLineNumber() to apply additional styles to the diff.
+ */
 export interface AnnotationContext {
+  /** Set by Gerrit and consumed by the plugin provided AddLayerFunc. */
+  readonly changeNum: number;
+  /** Set by Gerrit and consumed by the plugin provided AddLayerFunc. */
+  readonly path: string;
+  /** Set by Gerrit and consumed by the plugin provided AddLayerFunc. */
+  readonly line: GrDiffLine;
+  /** Set by Gerrit and consumed by the plugin provided AddLayerFunc. */
+  readonly contentEl: HTMLElement;
+  /** Set by Gerrit and consumed by the plugin provided AddLayerFunc. */
+  readonly lineNumberEl: HTMLElement;
+
   /**
-   * Method to add annotations to a content line.
+   * Can be called by the plugin to style a part of the given line of the
+   * context.
    *
    * @param offset The char offset where the update starts.
    * @param length The number of chars that the update covers.
@@ -56,7 +73,8 @@
   ): void;
 
   /**
-   * Method to add a CSS class to the line number TD element.
+   * Can be called by the plugin to style a part of the given line of the
+   * context.
    *
    * @param styleObject The style object for the range.
    * @param side The side of the update. ('left' or 'right')
@@ -66,23 +84,12 @@
 
 export interface AnnotationPluginApi {
   /**
-   * Register a function to call to apply annotations. Plugins should use
-   * GrAnnotationActionsContext.annotateRange and
-   * GrAnnotationActionsContext.annotateLineNumber to apply a CSS class to the
-   * line content or the line number.
-   *
-   * @param addLayerFunc The function
-   * that will be called when the AnnotationLayer is ready to annotate.
+   * Registers a callback for applying annotations. Gerrit will call the
+   * callback for every line of every file that is rendered and pass the
+   * information about the file and line as an AnnotationContext, which also
+   * provides methods for the plugin to style the content.
    */
-  addLayer(addLayerFunc: AddLayerFunc): AnnotationPluginApi;
-
-  /**
-   * The specified function will be called with a notify function for the plugin
-   * to call when it has all required data for annotation. Optional.
-   *
-   * @param notifyFunc See doc of the notify function below to see what it does.
-   */
-  addNotifier(notifyFunc: (n: NotifyFunc) => void): AnnotationPluginApi;
+  setLayer(callback: AnnotationCallback): AnnotationPluginApi;
 
   /**
    * The specified function will be called when a gr-diff component is built,
@@ -117,9 +124,10 @@
   ): AnnotationPluginApi;
 
   /**
-   * The notify function will call the listeners of all required annotation
-   * layers. Intended to be called by the plugin when all required data for
-   * annotation is available.
+   * For plugins notifying Gerrit about new annotations being ready to be
+   * applied for a certain range. Gerrit will then re-render the relevant lines
+   * of the diff and call back to the layer annotation function that was
+   * registered in addLayer().
    *
    * @param path The file path whose listeners should be notified.
    * @param start The line where the update starts.
diff --git a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-context.ts b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-context.ts
index 6a4da7b..9f9ba6a 100644
--- a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-context.ts
+++ b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-context.ts
@@ -29,14 +29,13 @@
  * @param lineNumberEl The TD element of the line number to
  * apply the annotation to using annotateLineNumber.
  * @param line The line object.
- * @param path The file path (eg: /COMMIT_MSG').
+ * @param path The file path (eg: '/COMMIT_MSG').
  * @param changeNum The Gerrit change number.
- * @param patchNum The Gerrit patch number.
  */
 export class GrAnnotationActionsContext implements AnnotationContext {
-  private _contentEl: HTMLElement;
+  contentEl: HTMLElement;
 
-  private _lineNumberEl: HTMLElement;
+  lineNumberEl: HTMLElement;
 
   line: GrDiffLine;
 
@@ -53,9 +52,8 @@
     path: string,
     changeNum: string | number
   ) {
-    this._contentEl = contentEl;
-    this._lineNumberEl = lineNumberEl;
-
+    this.contentEl = contentEl;
+    this.lineNumberEl = lineNumberEl;
     this.line = line;
     this.path = path;
     this.changeNum = Number(changeNum);
@@ -80,12 +78,12 @@
     styleObject: GrStyleObject,
     side: string
   ) {
-    if (this._contentEl?.getAttribute('data-side') === side) {
+    if (this.contentEl?.getAttribute('data-side') === side) {
       GrAnnotation.annotateElement(
-        this._contentEl,
+        this.contentEl,
         offset,
         length,
-        styleObject.getClassName(this._contentEl)
+        styleObject.getClassName(this.contentEl)
       );
     }
   }
@@ -97,8 +95,8 @@
    * @param side The side of the update. ('left' or 'right')
    */
   annotateLineNumber(styleObject: GrStyleObject, side: string) {
-    if (this._lineNumberEl?.classList.contains(side)) {
-      styleObject.apply(this._lineNumberEl);
+    if (this.lineNumberEl?.classList.contains(side)) {
+      styleObject.apply(this.lineNumberEl);
     }
   }
 }
diff --git a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api.ts b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api.ts
index 4abc6e1..95252cf 100644
--- a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api.ts
+++ b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api.ts
@@ -21,65 +21,37 @@
 import {EventType, PluginApi} from '../../../api/plugin';
 import {appContext} from '../../../services/app-context';
 import {
-  AddLayerFunc,
+  AnnotationCallback,
   AnnotationPluginApi,
   CoverageProvider,
-  NotifyFunc,
 } from '../../../api/annotation';
 
 export class GrAnnotationActionsInterface implements AnnotationPluginApi {
-  // Collect all annotation layers instantiated by getLayer. Will be used when
-  // notifying their listeners in the notify function.
+  /**
+   * Collect all annotation layers instantiated by createLayer. This is only
+   * used for being able to look up the appropriate layer when notify() is
+   * being called by plugins.
+   */
   private annotationLayers: AnnotationLayer[] = [];
 
-  private coverageProvider: CoverageProvider | null = null;
+  private coverageProvider?: CoverageProvider;
 
-  // Default impl is a no-op.
-  private addLayerFunc: AddLayerFunc = () => {};
+  private annotationCallback?: AnnotationCallback;
 
-  reporting = appContext.reportingService;
+  private readonly reporting = appContext.reportingService;
 
   constructor(private readonly plugin: PluginApi) {
-    // Return this instance when there is an annotatediff event.
     plugin.on(EventType.ANNOTATE_DIFF, this);
   }
 
-  /**
-   * Register a function to call to apply annotations. Plugins should use
-   * GrAnnotationActionsContext.annotateRange and
-   * GrAnnotationActionsContext.annotateLineNumber to apply a CSS class to the
-   * line content or the line number.
-   *
-   * @param addLayerFunc The function
-   * that will be called when the AnnotationLayer is ready to annotate.
-   */
-  addLayer(addLayerFunc: AddLayerFunc) {
-    this.addLayerFunc = addLayerFunc;
+  setLayer(annotationCallback: AnnotationCallback) {
+    if (this.annotationCallback) {
+      console.warn('Overwriting an existing plugin annotation layer.');
+    }
+    this.annotationCallback = annotationCallback;
     return this;
   }
 
-  /**
-   * The specified function will be called with a notify function for the plugin
-   * to call when it has all required data for annotation. Optional.
-   *
-   * @param notifyFunc See doc of the notify function below to see what it does.
-   */
-  addNotifier(notifyFunc: (n: NotifyFunc) => void) {
-    notifyFunc(
-      (path: string, startRange: number, endRange: number, side: Side) =>
-        this.notify(path, startRange, endRange, side)
-    );
-    return this;
-  }
-
-  /**
-   * The specified function will be called when a gr-diff component is built,
-   * and feeds the returned coverage data into the diff. Optional.
-   *
-   * Be sure to call this only once and only from one plugin. Multiple coverage
-   * providers are not supported. A second call will just overwrite the
-   * provider of the first call.
-   */
   setCoverageProvider(
     coverageProvider: CoverageProvider
   ): GrAnnotationActionsInterface {
@@ -98,23 +70,6 @@
     return this.coverageProvider;
   }
 
-  /**
-   * Returns a checkbox HTMLElement that can be used to toggle annotations
-   * on/off. The checkbox will be initially disabled. Plugins should enable it
-   * when data is ready and should add a click handler to toggle CSS on/off.
-   *
-   * Note1: Calling this method from multiple plugins will only work for the
-   * 1st call. It will print an error message for all subsequent calls
-   * and will not invoke their onAttached functions.
-   * Note2: This method will be deprecated and eventually removed when
-   * https://bugs.chromium.org/p/gerrit/issues/detail?id=8077 is
-   * implemented.
-   *
-   * @param checkboxLabel Will be used as the label for the checkbox.
-   * Optional. "Enable" is used if this is not specified.
-   * @param onAttached The function that will be called
-   * when the checkbox is attached to the page.
-   */
   enableToggleCheckbox(
     checkboxLabel: string,
     onAttached: (checkboxEl: Element | null) => void
@@ -148,16 +103,6 @@
     return this;
   }
 
-  /**
-   * The notify function will call the listeners of all required annotation
-   * layers. Intended to be called by the plugin when all required data for
-   * annotation is available.
-   *
-   * @param path The file path whose listeners should be notified.
-   * @param start The line where the update starts.
-   * @param end The line where the update ends.
-   * @param side The side of the update ('left' or 'right').
-   */
   notify(path: string, start: number, end: number, side: Side) {
     for (const annotationLayer of this.annotationLayers) {
       // Notify only the annotation layer that is associated with the specified
@@ -169,24 +114,25 @@
   }
 
   /**
-   * Should be called to register annotation layers by the framework. Not
-   * intended to be called by plugins.
+   * Factory method called by Gerrit for creating a DiffLayer for each diff that
+   * is rendered.
    *
-   * Don't forget to dispose layer.
-   *
-   * @param path The file path (eg: /COMMIT_MSG').
-   * @param changeNum The Gerrit change number.
+   * Don't forget to also call disposeLayer().
    */
-  getLayer(path: string, changeNum: number) {
+  createLayer(path: string, changeNum: number) {
+    if (!this.annotationCallback) return undefined;
     const annotationLayer = new AnnotationLayer(
       path,
       changeNum,
-      this.addLayerFunc
+      this.annotationCallback
     );
     this.annotationLayers.push(annotationLayer);
     return annotationLayer;
   }
 
+  /**
+   * Called by Gerrit for each diff renderer that had called createLayer().
+   */
   disposeLayer(path: string) {
     this.annotationLayers = this.annotationLayers.filter(
       annotationLayer => annotationLayer.path !== path
@@ -194,6 +140,10 @@
   }
 }
 
+/**
+ * An AnnotationLayer exists for each file that is being rendered. This class is
+ * not exposed to plugins, but being used by Gerrit's diff rendering.
+ */
 export class AnnotationLayer implements DiffLayer {
   private listeners: DiffLayerListener[] = [];
 
@@ -202,13 +152,13 @@
    *
    * @param path The file path (eg: /COMMIT_MSG').
    * @param changeNum The Gerrit change number.
-   * @param addLayerFunc The function
+   * @param annotationCallback The function
    * that will be called when the AnnotationLayer is ready to annotate.
    */
   constructor(
     readonly path: string,
     private readonly changeNum: number,
-    private readonly addLayerFunc: AddLayerFunc
+    private readonly annotationCallback: AnnotationCallback
   ) {
     this.listeners = [];
   }
@@ -230,7 +180,8 @@
   }
 
   /**
-   * Layer method to add annotations to a line.
+   * Called by Gerrit during diff rendering for each line. Delegates to the
+   * plugin provided callback for potentially annotating this line.
    *
    * @param contentEl The DIV.contentText element of the line
    * content to apply the annotation to using annotateRange.
@@ -243,18 +194,20 @@
     lineNumberEl: HTMLElement,
     line: GrDiffLine
   ) {
-    const annotationActionsContext = new GrAnnotationActionsContext(
+    const context = new GrAnnotationActionsContext(
       contentEl,
       lineNumberEl,
       line,
       this.path,
       this.changeNum
     );
-    this.addLayerFunc(annotationActionsContext);
+    this.annotationCallback(context);
   }
 
   /**
-   * Notify Layer listeners of changes to annotations.
+   * Notify layer listeners (which typically is just Gerrit's diff renderer) of
+   * changes to annotations after the diff rendering had already completed. This
+   * is indirectly called by plugins using the AnnotationPluginApi.notify().
    *
    * @param start The line where the update starts.
    * @param end The line where the update ends.
diff --git a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api_test.js b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api_test.js
index 7ae34cf..9811f99 100644
--- a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api_test.js
+++ b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-annotation-actions-js-api_test.js
@@ -59,9 +59,9 @@
       assert.equal(context.line, line);
       assert.equal(context.changeNum, changeNum);
     };
-    annotationActions.addLayer(testLayerFunc);
+    annotationActions.setLayer(testLayerFunc);
 
-    const annotationLayer = annotationActions.getLayer(
+    const annotationLayer = annotationActions.createLayer(
         '/dummy/path', changeNum);
 
     const lineNumberEl = document.createElement('td');
@@ -72,27 +72,19 @@
   test('add notifier', () => {
     const path1 = '/dummy/path1';
     const path2 = '/dummy/path2';
-    const annotationLayer1 = annotationActions.getLayer(path1, 1);
-    const annotationLayer2 = annotationActions.getLayer(path2, 1);
+    annotationActions.setLayer(context => {});
+    const annotationLayer1 = annotationActions.createLayer(path1, 1);
+    const annotationLayer2 = annotationActions.createLayer(path2, 1);
     const layer1Spy = sinon.spy(annotationLayer1, 'notifyListeners');
     const layer2Spy = sinon.spy(annotationLayer2, 'notifyListeners');
 
-    let notify;
-    let notifyFuncCalled;
-    const notifyFunc = n => {
-      notifyFuncCalled = true;
-      notify = n;
-    };
-    annotationActions.addNotifier(notifyFunc);
-    assert.isTrue(notifyFuncCalled);
-
     // Assert that no layers are invoked with a different path.
-    notify('/dummy/path3', 0, 10, 'right');
+    annotationActions.notify('/dummy/path3', 0, 10, 'right');
     assert.isFalse(layer1Spy.called);
     assert.isFalse(layer2Spy.called);
 
     // Assert that only the 1st layer is invoked with path1.
-    notify(path1, 0, 10, 'right');
+    annotationActions.notify(path1, 0, 10, 'right');
     assert.isTrue(layer1Spy.called);
     assert.isFalse(layer2Spy.called);
 
@@ -101,7 +93,7 @@
     layer2Spy.resetHistory();
 
     // Assert that only the 2nd layer is invoked with path2.
-    notify(path2, 0, 20, 'left');
+    annotationActions.notify(path2, 0, 20, 'left');
     assert.isFalse(layer1Spy.called);
     assert.isTrue(layer2Spy.called);
   });
@@ -143,7 +135,8 @@
   });
 
   test('layer notify listeners', () => {
-    const annotationLayer = annotationActions.getLayer('/dummy/path', 1);
+    annotationActions.setLayer(context => {});
+    const annotationLayer = annotationActions.createLayer('/dummy/path', 1);
     let listenerCalledTimes = 0;
     const startRange = 10;
     const endRange = 20;
diff --git a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-js-api-interface-element.ts b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-js-api-interface-element.ts
index 830fb92..8689ad2 100644
--- a/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-js-api-interface-element.ts
+++ b/polygerrit-ui/app/elements/shared/gr-js-api-interface/gr-js-api-interface-element.ts
@@ -252,8 +252,8 @@
     for (const cb of this._getEventCallbacks(EventType.ANNOTATE_DIFF)) {
       const annotationApi = (cb as unknown) as GrAnnotationActionsInterface;
       try {
-        const layer = annotationApi.getLayer(path, changeNum);
-        layers.push(layer);
+        const layer = annotationApi.createLayer(path, changeNum);
+        if (layer) layers.push(layer);
       } catch (err) {
         this.reporting.error(err);
       }
diff --git a/polygerrit-ui/app/samples/coverage-plugin.js b/polygerrit-ui/app/samples/coverage-plugin.js
index 9b2b687..8d321c7 100644
--- a/polygerrit-ui/app/samples/coverage-plugin.js
+++ b/polygerrit-ui/app/samples/coverage-plugin.js
@@ -41,7 +41,7 @@
   const coverageStyle = styleApi.css('background-color: #EF9B9B !important');
   const emptyStyle = styleApi.css('');
 
-  annotationApi.addLayer(context => {
+  annotationApi.setLayer(context => {
     if (Object.keys(coverageData).length === 0) {
       // Coverage data is not ready yet.
       return;
@@ -64,19 +64,13 @@
       }
     }
   }).enableToggleCheckbox('Display Coverage', checkbox => {
-    // Checkbox is attached so now add the notifier that will be controlled
-    // by the checkbox.
-    // Checkbox will only be added to the file diff page, in the top right
-    // section near the "Diff view".
-    annotationApi.addNotifier(notifyFunc => {
-      populateWithDummyData(coverageData);
-      checkbox.disabled = false;
-      checkbox.onclick = e => {
-        displayCoverage = e.target.checked;
-        Object.keys(coverageData).forEach(file => {
-          notifyFunc(file, 0, coverageData[file].totalLines, 'right');
-        });
-      };
-    });
+    populateWithDummyData(coverageData);
+    checkbox.disabled = false;
+    checkbox.onclick = e => {
+      displayCoverage = e.target.checked;
+      Object.keys(coverageData).forEach(file => {
+        annotationApi.notify(file, 0, coverageData[file].totalLines, 'right');
+      });
+    };
   });
-});
\ No newline at end of file
+});