Do I need to use storage gaps when importing upgradeable contracts?

Hello guys,

I would like to share with you my implementation and I would like to know if my assumption about the memory layout is correct.

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.24;

import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeable.sol";
import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";

/// @title Content Distributor contract.
/// @notice Use this contract to handle all needed logic for distributors.
/// @dev This contract inherits from Ownable and ERC165, and implements the IDistributor interface.
/// Extending upgradeable contracts in a non-upgradeable contract to extend ERC-7201: Namespaced Storage Layout
/// Same as below with __gap the issue could happen using this contract as implementation and receiving delegated calls.
/// This contract can be deployed without needing to upgrade.
contract Distributor is
    Initializable,
    ERC165Upgradeable,
    OwnableUpgradeable,
{
    
    /// @notice The URL to the distribution.
    /// Since this is a contract considered as implementation for beacon proxy,
    /// we need to reserve a gap for endpoint to avoid memory layout getting mixed up.
    /// https://docs.openzeppelin.com/upgrades-plugins/1.x/writing-upgradeable
    string private endpoint;

    /// @notice We use this method to initialize store from "BeaconProxy"
    function initialize(
        string memory _endpoint,
        address _owner
    ) public initializer {
        __ERC165_init();
        __Ownable_init(_owner);
        
        if (bytes(_endpoint).length == 0) revert InvalidEndpoint();
        endpoint = _endpoint;
    }

    .....

    // Reserved space for future storage variables to prevent storage conflicts
    uint256[20] private __gap;
}

The factory

// SPDX-License-Identifier: MIT
// NatSpec format convention - https://docs.soliditylang.org/en/v0.5.10/natspec-format.html
pragma solidity ^0.8.24;

import "@openzeppelin/contracts/utils/Pausable.sol";
import "@openzeppelin/contracts/proxy/beacon/BeaconProxy.sol";
import "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol";
import "./Distributor.sol";


// Each distributor has their own contract. The problem with this approach is that each contract
// has its own implementation. If in the future we need to improve the distributor contract,
// we can't deploy and upgrade each contract individually to update the implementation.
// Even worse, if the contract is not upgradeable, the implementation cannot be updated,
// requiring a new deployment, which is a significant hassle.

// The solution involves using a beacon proxy pattern:
// beaconProxy -> beacon
//             -> beacon
//             -> beacon -> implementation
//             -> beacon
//             -> beacon 

contract DistributorFactory is UpgradeableBeacon, Pausable {

    constructor(
        address implementation,
        address initialOwner
    ) UpgradeableBeacon(implementation, initialOwner) Pausable() {}

    /// @notice Function to pause the contract, preventing the creation of new distributors.
    /// @dev Can only be called by the owner of the contract.
    function pause() external onlyOwner {
        _pause();
    }

    /// @notice Function to unpause the contract, allowing the creation of new distributors.
    /// @dev Can only be called by the owner of the contract.
    function unpause() external onlyOwner {
        _unpause();
    }

    function register(string calldata _endpoint) external whenNotPaused {
        // not allowed duplicated endpoints
        if (registry[_endpoint] != address(0))
            revert DistributorAlreadyRegistered();

        // initialize storage layout from Distributor contract..
        bytes memory data = abi.encodeWithSignature(
            "initialize(string,address)",
            _endpoint, _msgSender()
        );

        address newContract = address(new BeaconProxy(address(this), data));
        registry[_endpoint] = _msgSender();
    }
}

Best Regards..

@Geolffrey_Mena Can you elaborate on what is your assumption or question?

Hey @ericglau. If you have a implementation contract in beacon pattern, the same issue described here: https://docs.openzeppelin.com/upgrades-plugins/1.x/writing-upgradeable with memory layouts will happen if you dont be mindful with storage collision right? In this example we inherit from Upgradeable contracts to extend ERC-7201 features (the contract itself is not upgradeable) and created a gap in implementation contract to tackle it. Just wondering if this is something needed or not in case that Openzeppellin is doing something underneath?

Storage gaps and ERC-7201 namespaced storage layout are two different ways of allowing base contracts' storage variables to be added to in the future without affecting child contracts which inherit it. ERC-7201 (if all of your contracts are using ERC-7201 instead of storage gaps) additionally has the benefit of allowing inheritance order of base contracts to be changed, without affecting the resulting storage layout.

If you don't intend for your Distributor to be inherited as a base contract from other contracts, then it is not necessary to include a storage gap. But it is fine to do so.

In our case, the Distributor contract is the implementation in the UpgradeableBeacon, so the calls it receives will be delegated through the BeaconProxy. If this is the case, is it possible that the implementation, when undergoing changes (such as adding attributes) without proper care in storage, might experience collisions? To address this, we inherit from upgradeable contracts only to avoid having to rewrite the namespaces of OwnableUpgradeable, which are already ERC-7201. The Distributor contract is not upgradeable itself. Regarding the gap, we leave it for the proper use of the Distributor contract in case modifications to the implementation arise.

If you are using the UpgradeableBeacon to change your implementation to a new version of Distributor in the future, then yes, you need to ensure the storage layout remains compatible in the new versions.