Initial state once cell value

[IpFruion]

Proposal

Problem statement

Provide a mechanism to lazily load a value from an initial state only when successful.

Motivating examples or use cases

Currently, OnceCell<T> and LazyCell<T> provide mechanisms to defer initialization until later, however they do not seem to provide a way to solve handling the full problem statement above.

Example 1 LazyCell<T>

Given some structure that is a long running structure in an application and it's new and request function

pub struct App {
   client: LazyCell<Result<Client, ClientError>>,
}
impl App {
   pub fn new(settings: ClientSettings) -> Self {
      App { 
         client: LazyCell::new(move || {
            Ok(Client::new(settings)?)
         })
      }
   }
   pub fn request(self, ...) -> Result<Response, RequestError> {
      // Initialization of the client only attempts once
      // This can lead to issues where something on the system might have changed and retrying to build the client might be necessary
      let response = self.client.map_err(|err| err.clone())?.send(...)?;
      Ok(response)
   }
}

We can see how we have a way to initialize this client but it can error and this function is only run once (in the request(...) function). A developer could need a mechanism to instead of only being able to attempt to initialize a value once, but to handle attempting again if the function fails (like in OnceCell::get_or_try_init where the value doesn't get initialized if it failed to init.

Example 2 OnceCell<T>

To attempt to fix being able to attempt to initialize the client (and only initialize on success), we could instead use OnceCell<T> to create this scenario

pub struct App {
   settings: ClientSettings,
   client: OnceCell<Client>
}
impl App {
   pub fn new(settings: ClientSettings) -> Self {
      App {
         settings,
         client: OnceCell::new(),
      }
   }
   pub fn request(self, ...) -> Result<Response, RequestError> {
      // we have to clone these settings here every time, even if the client is already initialized
      let settings = self.settings.clone();
      let response = self.client.get_or_try_init(move || {
         let client = Client::new(settings)?;
         Ok(client)
      })?.send(...)?;
      Ok(response)
   }
}

This gives us the ability to attempt initialization but keeps a long lasting ClientSettings around that is no longer going to be used after the Client is initialized. These ClientSettings could be de-allocated with a different structure.

Solution sketch

The thought process between these solutions is to combine both the worlds of LazyCell and OnceCell together to allow for some initial state to be able to be passed along and used in the attempts for initialization. Let's call this new structure StateCell (not set on the name)

pub struct StateCell<I, T> {
    state: UnsafeCell<State<I, T>>,
}

// Very similar to `LazyCell<T>`'s internal `State<T, F>` 
enum State<I, T> {
   Uninit(I),
   Init(T),
   Poisoned,
}

This internally looks almost exactly like LazyCell<T>'s internal structure, except without the bounds that LazyCell<T, F = fn() -> T> meaning the State::Uninit is by default a function.

We can then use the above to construct some of the OnceCell<T> functions to get the get_or_try_init functionality.

impl StateCell<I, T> {
    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
    where
        F: FnOnce(&I) -> Result<T, E>,
    {
        // SAFETY:
        // This invalidates any mutable references to the data. The resulting
        // reference lives either until the end of the borrow of `self` (in the
        // initialized case) or is invalidated in `really_init` (in the
        // uninitialized case; `really_init` will create and return a fresh reference).
        let state = unsafe { &*self.state.get() };
        match state {
            State::Init(i) => Ok(i),
            State::Poisoned => panic_poisoned(),
            // SAFETY: The state is uninitialized.
            State::Uninit(_) => unsafe { self.really_try_init(f) },
        }
    }
    unsafe fn really_try_init<F, E>(&self, f: F) -> Result<&T, E>
    where
        F: FnOnce(&I) -> Result<T, E>,
    {
        // Swap out the initial state with the poisoned state to make sure the value is only alive for this function call
        // This helps prevent a case where `f` panics and the state should be poisoned like in `LazyCell`
        let i = unsafe { self.swap_poisoned() };
        let data = match f(&i) {
            Ok(data) => data,
            Err(err) => {
                // An error occured and thus we want to reset the internal state
                unsafe { self.state.get().write(State::Uninit(i)) };
                return Err(err);
            }
        };
        
        // Otherwise we have been successful and we want to initialize the data.
        Ok(unsafe { self.set_data(data) })
    }
}

The content of these functions are very similar to LazyCell however they provide this functionality later when the function is called to get access.

Example Usage

Our App example can now be different

pub struct App {
   client: StateCell<ClientSettings, Client>
}
impl App {
   pub fn new(settings: ClientSettings) -> Self {
      App {
         settings,
         client: StateCell::new(settings),
      }
   }
   pub fn request(self, ...) -> Result<Response, RequestError> {
      let response = self.client.get_or_try_init(|settings: &ClientSettings| {
         let client = Client::new(settings)?;
         Ok(client)
      })?.send(...)?;
      Ok(response)
   }
}

Alternatives

One alternative could be to relax the bounds on F in LazyCell<T, F> i.e. having the F function produce F: Fn() -> Result<T, E> but this has some compexity with the trait bounds in place for deref and other implementations.

Another option would be to completely relax on LazyCell<T, F> to allow F to not be a function but instead some initial value and provide some of the functions of OnceCell onto LazyCell only when F is not a FnOnce function but this goes against a lot of the initial design of LazyCell which seems like it would break some of this.

The reason I didn't recommend these as solutions is because it would produce a significant change to the API that is already there instead of adding a new solution as discussed above.

This can definitely be done as a crate on crates.io (I am new to the process of open source development and wanted to bring this problem, and possible solution, to where I think it would serve the wider audience as a whole). Another reason would be since this does have some unsafe work in it (very similar to the LazyCell and OnceCell safety guarantees), it would be harder to get public view into the solution since I would be wary of adding a singular crate for this structure that seems to fit in-between two standard library structures.

Links and related work

I don't have any other links here yet but will add them when / if I find them. I have started drafting up what the solution might look like in code (plus some tests to see if I am on the right track) and can add the fork (branch? not sure what the best way to contribute would be) to this issue.

Lastly, I would just like to state that I am new to the contributing space and I am happy to help discuss, chat, and design out the solution (or other solutions that others might recommend). I may be missing something obvious that already solves this problem and would be glad to review it if it meets the needs I have in mind. I appreciate any feedback and suggestions.

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[joshtriplett]

We talked about this in today's @rust-lang/libs-api meeting.

1. We'd like to see State made a public type, so that people can use State directly without a cell when they can use &mut self and don't need interior mutability. (Would that work for your use case, or do you need interior mutability?) State should then have appropriate functions directly.

2. We'd like to see the poisoning separated into a different type. The work on non-poisoning mutexes has suggested a Poison<T> type, which would be the same type used to implement a poisoning mutex atop a non-poisoning mutex. In this case, State<I, Poison<T>> would work for cases that need poisoning, and some use cases would be able to use State<I, T> because they don't need poisoning (e.g. an infallible initialization function that can take &I). Functionality that only works with poisoning can require State<I, Poison<T>>.

3. For the cell version, we felt that we'd prefer the version proposed in Allow arbitrary initializers in LazyCell and LazyLock #670 rather than here: reusing the LazyCell/LazyLock types with relaxed bounds on the "function", rather than creating a new StateCell (and presumably a thread-safe StateLock). Would that work for your use case?

[IpFruion]

Thanks for the quick response @joshtriplett and the consideration!

We talked about this in today's @rust-lang/libs-api meeting.

1. We'd like to see `State` made a public type, so that people _can_ use `State` directly without a cell when they can use `&mut self` and don't need interior mutability. (Would that work for your use case, or do you _need_ interior mutability?) `State` should then have appropriate functions directly.

For the use case above I think interior mutability would be necessary since I would want the value to be initialized "Once" when successful. However, I do see how one could just leverage mutability to get a similar job done without the need for a cell via exposing State. Taking the example case, I would probably want to cut down on the number of times I would need to lock the App value to gain mutable access to initialize the client. So this could be introduced by some RwLock where it attempts to see if the client is initialized first by some read operation and then if that fails to perform the write operation to update it. In general, I would probably want some form of interior mutability that allows this to happen though.

2. We'd like to see the poisoning separated into a different type. The work on non-poisoning mutexes has suggested a `Poison<T>` type, which would be the same type used to implement a poisoning mutex atop a non-poisoning mutex. In this case, `State<I, Poison<T>>` would work for cases that need poisoning, and some use cases would be able to use `State<I, T>` because they don't need poisoning (e.g. an infallible initialization function that can take `&I`). Functionality that only works with poisoning can require `State<I, Poison<T>>`.

Hopefully I am understanding this concept correctly. I like this design because this helps make the StateCell more versatile however in the cases that I have seen, I am not sure where it wouldn't be poisoned. As a demo of "what could be" here is what I am understanding:

struct State<I, T> {
   Uninit(I)
   Init(T)
}

impl StateCell<I, T> {
   pub fn set(&self, value: T) -> Result<&T, T> {
      ...
   }
}

impl StateCell<I, Poison<T>> {
    pub fn get_or_init<F: FnOnce(I) -> T>(&self, f: F) -> Option<&T> {
      ...
    }
   pub fn get_or_try_init<E, F: FnOnce(&I) -> Result<T, E>>(&self, f: F) -> Result<&T, E> {
      ...
    }
}

The caveat being, the type exposure would be a little odd i.e. when creating a type it would have to include being able to be poisoned StateCell<Client, Poison<ClientSettings>> (which I guess makes sense if you want to use those functions).

3. For the cell version, we felt that we'd prefer the version proposed in [Allow arbitrary initializers in `LazyCell` and `LazyLock` #670](https://github.com/rust-lang/libs-team/issues/670) rather than here: reusing the `LazyCell`/`LazyLock` types with relaxed bounds on the "function", rather than creating a new `StateCell` (and presumably a thread-safe `StateLock`). Would that work for your use case?

I like the idea to relax the bounds on the "function" rather than creating a new structure here (I had this in my alternatives section and glad to see there is other design ideas on that 😄 ) but I am not sure it works for my use case because of being able to "attempt to initialize" that the OnceCell provides with get_or_try_init.

Example pulled from the issue above

Lets say the Bar::parse function is instead:

impl Bar {
   pub fn parse(raw: &[u8], foo: &Foo) -> Result<Bar, ParseError>;
}

By the looks of these changes there is not a force_with_try functionality which would attempt to initialize if the Result<Bar, ParseError> was true i.e.

impl LazyCell<T, F> {
  pub fn force_with_try<E, I: FnOnce(&F) -> Result<T, E>>(this: &Self, init: I) -> Result<&T, E>;
}

If this was part of that design, it would satisfy the use case I am envisioning and would allow for this lazy initialization attempt using a stored value when created.

[programmerjake]

I think the idea is that you don't need Poison<T> for get_or_[try_]init functions when they take &T since if the function panics or errors it can still leave the T there since it's presumably unmodified. you'd only need Poison<T> when taking T by value since if the function panics or errors you already moved the T so you don't have a T to put back there anymore, hence why it'll get set to Poison::Poisoned (or whatever that variant's called).

so you could have (names to be debated):

impl<I, T> StateCell<I, T> {
    pub fn get_or_init_ref(&self, f: impl FnOnce(&T) -> I) -> &I { ... }
    pub fn get_or_try_init_ref<E>(&self, f: impl FnOnce(&T) -> Result<I, E>) -> Result<&I, E> { ... }
}

impl<I, T> StateCell<I, Poison<T>> {
    /// poisons self on panic
    pub fn get_or_init_move(&self, f: impl FnOnce(T) -> I) -> &I { ... }
    /// poisons self on panic or error
    pub fn get_or_try_init_move<E>(&self, f: impl FnOnce(T) -> Result<I, E>) -> Result<&I, E> { ... }
}

[programmerjake]

or maybe even:

impl<I, T> StateCell<I, Poison<T>> {
    /// poisons self on panic, but not on error
    pub fn get_or_try_init_move2<E>(&self, f: impl FnOnce(T) -> Result<I, (T, E)>) -> Result<&I, E> { ... }
}

[IpFruion]

Oh yeah that makes total sense @programmerjake. I will say with one caveat on probably not having a

impl<I, T> StateCell<I, Poison<T>>
   /// poisons self on panic or error
   pub fn get_or_try_init_move<E>(&self, f: impl FnOnce(T) -> Result<I, E>) -> Result<&I, E> { ... }
}

function due to it poisoning the cell when a possibly intentional error occurs 🤷

I also thought about the

impl<I, T> StateCell<I, Poison<T>> {
    /// poisons self on panic, but not on error
    pub fn get_or_try_init_move2<E>(&self, f: impl FnOnce(T) -> Result<I, (T, E)>) -> Result<&I, E> { ... }
}

functionality which could provide the _move functionality better if desired while keeping the _ref impls alone.

To sum up I like these bounds (naming TBD):

pub struct StateCell<I, T> {
    state: UnsafeCell<State<I, T>>,
}

enum State<I, T> {
   Uninit(I),
   Init(T),
}

impl<I, T> StateCell<I, T> {
   pub fn set(&self, value: T) -> Result<&T, T> { ... }
   pub fn get(&self) -> Option<&T> { ... }
   pub fn get_or_init_ref<F: FnOnce(&I) -> T>(&self, f: F) -> &T { ... }
   pub fn get_or_try_init_ref<E: F: FnOnce(&I) -> Result<T, E>>(&self, f: F) -> Result<&T, E> { ... }
}

impl<I, T> StateCell<I, Poison<T>> {
   pub fn get_or_init<F: FnOnce(I) -> T>(&self, f: F) -> &T { ... }
   pub fn get_or_try_init<E, F: FnOnce(I) -> Result<T, (I, E)>>(&self, f: F) -> Result<&T, E> { ... }
}

[IpFruion]

With the above design however (with the Poison<T> bounds), some functionality might be slightly more complex from the developer perspective i.e. the get function.

pub fn main() {
   let state: StateCell<u32, Poison<String>> = StateCell::new(43);
   ...
   // Where it has to check for if the value has been poisoned or not
   // We could offer some helper function for `StateCell<I, Poison<T>>` that flattens this to an `Option<&T>`
   if let Some(Poison::Value(v)) = state.get() {
      ...
   }
}

It also allows for the set function to manually be able to poison a value, which I am not opposed too but might have a little different impact i.e. state.set(Poison::Poisoned)

[programmerjake]

I was assuming the Poison would be in the not-yet-initialized value, so get and set don't have to deal with poisoning, get just returns None and set just takes the final value without Poison.

[IpFruion]

I was assuming the Poison would be in the not-yet-initialized value, so get and set don't have to deal with poisoning, get just returns None and set just takes the final value without Poison.

This wouldn't be the case with the get and set in the State<I, T> impl since that would allow T here to be Poison<T> so if one wanted to use get with the State<I, Poison<T>> then it would return Option<Poison<T>> and similar for the set function

[programmerjake]

I was assuming the Poison would be in the not-yet-initialized value, so get and set don't have to deal with poisoning, get just returns None and set just takes the final value without Poison.

This wouldn't be the case with the get and set in the State<I, T> impl since that would allow T here to be Poison<T> so if one wanted to use get with the State<I, Poison<T>> then it would return Option<Poison<T>> and similar for the set function

I was thinking State::<A, Poison<B>>::get() returns &A, not &Poison<B>. imo having the actual type that get returns come first (like LazyCell) and the init-time state come second makes more sense to me, since the type get() returns is the more important generic parameter.

[IpFruion]

@programmerjake I see so get would return the uninit state yeah I was more thinking about the other OnceCell<T>::get where it returns the reference if the value is set

I agree though which is why I am a little confused as to the necessity of an extra Poison<T> type other than marking some of that functionality as not able to be poisoned. Either way I am good with including it in the design or not

[programmerjake]

I'm going to write what I meant out in excessive detail since you seem confused about it:

While writing this out I realized that with just UnsafeCell<State<_, _>> it is unsound to pass references to ArgForInitialize to the init function because you can call get_or_init_ref recursively so the inner replaces the state while the &ArgForInitialize from the outer call is still accessible.

pub enum State<Initalized, ArgForInitialize> {
    Uninit(ArgForInitialize),
    Init(Initialized),
}

pub struct StateCell<Initalized, ArgForInitialize>(UnsafeCell<State<Initalized, ArgForInitialize>>);

impl<Initalized, ArgForInitialize> StateCell<Initalized, ArgForInitialize> {
    pub fn get(&self) -> Option<&Initialized> {
        // Safety: once self.0 is Init, it's not changed without mutable access to self which causes the returned lifetime to be expired
        unsafe {
            match *self.0.get() {
                State::Init(ref retval) => Some(retval),
                _ => None,
            }
        }
    }
    // *unsound*
    pub fn get_or_init_ref(&self, f: impl FnOnce(&ArgForInitialize) -> Initialized) -> &Initialized {
        ...
    }
}

impl<Initalized, ArgForInitialize> State<Initalized, Poison<ArgForInitialize>> {
    pub fn get_or_init(&self, f: impl FnOnce(ArgForInitialize) -> Initalized) -> &Initalized { ... }
}

[programmerjake]

so, since passing a reference to the initialization function is unsound due to recursion without some extra state to prevent recursion, I'm nominating this for some redesign of the suggested API.
@rustbot label +I-libs-api-nominated