extensions/browser/api/socket/write_quota_checker.h

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef EXTENSIONS_BROWSER_API_SOCKET_WRITE_QUOTA_CHECKER_H_
#define EXTENSIONS_BROWSER_API_SOCKET_WRITE_QUOTA_CHECKER_H_

#include <stddef.h>

#include <map>

#include "base/auto_reset.h"
#include "extensions/browser/browser_context_keyed_api_factory.h"
#include "extensions/common/extension_id.h"

namespace extensions {

// Tracks bytes quota per extension.
class WriteQuotaChecker : public BrowserContextKeyedAPI {
 public:
  class ScopedBytesLimitForTest {
   public:
    ScopedBytesLimitForTest(WriteQuotaChecker* checker, size_t new_bytes_limit);
    ~ScopedBytesLimitForTest();

   private:
    base::AutoReset<size_t> scoped_bytes_limit_;
  };

  static WriteQuotaChecker* Get(content::BrowserContext* context);

  // BrowserContextKeyedAPI:
  static BrowserContextKeyedAPIFactory<WriteQuotaChecker>* GetFactoryInstance();
  static const char* service_name() { return "WriteQuotaChecker"; }

  explicit WriteQuotaChecker(content::BrowserContext* context);
  ~WriteQuotaChecker() override;

  // Attempts to take more bytes to write. Returns false if the attempt would
  // exceed the limit.
  bool TakeBytes(const ExtensionId& extension_id, size_t bytes);

  // Puts bytes back to the pool after done with writing.
  void ReturnBytes(const ExtensionId& extension_id, size_t bytes);

 private:
  friend class BrowserContextKeyedAPIFactory<WriteQuotaChecker>;

  // Max pending write bytes.
  size_t bytes_limit_ = 0;

  // Tracked writing bytes per extension.
  std::map<ExtensionId, size_t> bytes_used_map_;
};

}  // namespace extensions

#endif  // EXTENSIONS_BROWSER_API_SOCKET_WRITE_QUOTA_CHECKER_H_