Home > database >  Storing different implementations of a trait in a HashMap
Storing different implementations of a trait in a HashMap

Time:10-26

I've come up against an interesting issue. Here is a playground link I'm trying to write a flexible task executor that is driven by providing JSON objects such as the following.

{
    "task_type": "foobar",
    "payload": {}
}

The task executor would use the task_type to determine which executor to call and then pass the associated payload. Most of my types support generics because different task_types would support different payloads. The way I've chosen to solve this problem is by creating a trait Task and a single implementor (for now) called RegisteredTask which essentially just calls a function pointer.

My only compilation problem is the following which indicates that I would need to restrict my registry to supporting implementations of Task that take the same inputs and return the same outputs (bye bye generics).

Compiling playground v0.0.1 (/playground)
error[E0191]: the value of the associated types `Input` (from trait `Task`), `Output` (from trait `Task`) must be specified
  --> src/lib.rs:30:36
   |
5  |     type Input;
   |     ----------- `Input` defined here
6  |     type Output;
   |     ------------ `Output` defined here
...
30 |     tasks: HashMap<String, Box<dyn Task>>,
   |                                    ^^^^ help: specify the associated types: `Task<Input = Type, Output = Type>`

After watching this video I understand the the compiler needs to know the sizes of everything at compile time, but I thought that by putting the trait in a Box I would be able to avoid that.

What is the root cause of my problem, and is the idiomatic way around this?

use std::collections::HashMap;

/// Task defines something that can be executed with an input and respond with an output
trait Task {
    type Input;
    type Output;
    fn execute(&self, input: Self::Input) -> Self::Output;
}

/// RegisteredTask is an implementation of Task that simply executes a predefined function and
/// returns it's output. It is generic so that there can be many different types of RegisteredTasks
/// that do different things.
struct RegisteredTask<T, K> {
    action: fn(T) -> K,
}

impl<T, K> Task for RegisteredTask<T, K> {
    type Input = T;
    type Output = K;

    /// Executes the registered task's action function and returns the output
    fn execute(&self, input: Self::Input) -> Self::Output {
        (self.action)(input)
    }
}

/// Maintains a collection of anything that implements Task and associates them with a name. This
/// allows us to easily get at a specific task by name.
struct Registry {
    tasks: HashMap<String, Box<dyn Task>>,
}

impl Registry {
    /// Registers an implementation of Task by name
    fn register<T, K>(&mut self, name: &str, task: Box<dyn Task<Input = T, Output = K>>) {
        self.tasks.insert(name.to_string(), task);
    }

    /// Gets an implementation of task by name if it exists, over a generic input and output
    fn get<T, K>(&self, name: &str) -> Option<&Box<dyn Task<Input = T, Output = K>>> {
        self.tasks.get(name)
    }
}

/// An example input for a registered task
#[derive(Debug)]
pub struct FooPayload {}

/// An example output for a registered task
#[derive(Debug)]
pub struct FooOutput {}

/// Executes the named task from the registry and decodes the required payload from a JSON string. The
/// decoding is not shown here for simplicity.
fn execute_task(registry: &Registry, name: &str, json_payload: String) {
    match name {
        "foobar" => {
            let task = registry.get::<FooPayload, FooOutput>(name).unwrap();
            let output = task.execute(FooPayload {});
            println!("{:?}", output)
        }
        _ => {}
    }
}

#[test]
fn test_execute_task() {
    // Create a new empty registry
    let mut registry = Registry {
        tasks: HashMap::new(),
    };

    // Register an implementation of Task (RegisteredTask) with the name "foobar"
    registry.register::<FooPayload, FooOutput>(
        "foobar",
        Box::new(RegisteredTask::<FooPayload, FooOutput> {
            // This is the function that should be called when this implementation of Task is invoked
            action: |payload| {
                println!("Got payload {:?}", payload);
                FooOutput {}
            },
        }),
    );

    // Attempt to invoke whatever implementation of Task (in this case a RegisteredTask) is named 'foobar'.
    execute_task(&registry, "foobar", String::from("{}"));
}

CodePudding user response:

The reason the compiler is upset with you is that generic types with different generic value are actually different types. This means that if you have a trait Foo<T>, and use it with T == Bar or T == Barcode, you have two different traits internally, not one. The same goes with associated types. A struct that implements Foo<InputType = Bar> and different struct that implements Foo<InputType = Barcode> implement different traits.

That is why writing dyn Task in your case does not make sense. Compiler needs to know which exact variant of Task trait objects you want to store.

Having generic parameters (or associated types) as method arguments can also prove to be a challenge later, as typically traits with such traits are not object safe, so you cannot use dyn Task for them.

It's also not quite clear how you would use the result of tasks, as you don't know specific types that the trait objects from the registry would return.

Better design would be to have just one non-generic trait. Then each struct implementing this trait could be constructed with Payload type as input and store it internally. Then .execute() can be called without parameter.

What type to return as the result of .execute() depends on the way you are going to use it. It could be an enum or serialized value.

Probably there is a better design for what you are trying to achieve, but this is what comes to my mind.

CodePudding user response:

Thanks to @Maxim, I redesigned my implementation and instead of having a generic registered task type, I created non-generic task-specific types. For example, here is a task that executes a shell command.

pub struct Exec {
    command: String,
    args: Vec<String>,
}

I went ahead and implemented the Task trait for this task, and because the payloads are now stored type specifically on the Task implementations, I don't have to worry about any generic implementations of the trait.

impl Task for Exec {
    async fn execute(&self) -> Result<Value> {
        // ....
    }
}
  • Related