1 /* 2 * Copyright (C) 2023 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #pragma once 18 19 #include <jni.h> 20 21 #include <memory> 22 23 #include "absl/status/statusor.h" 24 #include "fcp/client/simple_task_environment.h" 25 #include "fcp/jni/jni_util.h" 26 #include "fcp/protos/plan.pb.h" 27 28 /** 29 * A wrapper around a Java class for callbacks into Java. This is required for 30 * functionality we do not provide in c++ code, such as Clearcut logging, 31 * checking for device conditions, or querying apps for examples. 32 */ 33 34 namespace fcp { 35 namespace client { 36 namespace engine { 37 namespace jni { 38 39 using ::google::internal::federated::plan::ExampleSelector; 40 41 // Descriptions of Java classes we call into - full class path, and method 42 // names + JNI signatures. 43 struct SimpleTaskEnvironmentImplClassDesc { 44 static constexpr fcp::jni::JavaMethodSig kTrainingConditionsSatisfied = { 45 "trainingConditionsSatisfied", "()Z"}; 46 static constexpr fcp::jni::JavaMethodSig kCreateExampleIterator = { 47 "createExampleIteratorWithContext", 48 "([B[B)Lcom/android/federatedcompute/services/training/jni/" 49 "JavaExampleIterator;"}; 50 }; 51 52 class SimpleTaskEnvironmentWrapperImpl : public SimpleTaskEnvironment { 53 public: 54 explicit SimpleTaskEnvironmentWrapperImpl(JavaVM *jvm, 55 jobject simple_task_env); 56 57 ~SimpleTaskEnvironmentWrapperImpl() override; 58 59 std::string GetBaseDir() override; 60 61 std::string GetCacheDir() override; 62 63 bool TrainingConditionsSatisfied() override; 64 65 absl::StatusOr<std::unique_ptr<ExampleIterator>> CreateExampleIterator( 66 const google::internal::federated::plan::ExampleSelector 67 &example_selector) override; 68 69 absl::StatusOr<std::unique_ptr<ExampleIterator>> CreateExampleIterator( 70 const ExampleSelector &example_selector, 71 const SelectorContext &selector_context) override; 72 73 private: 74 JavaVM *jvm_; 75 jobject jthis_; 76 jmethodID training_conditions_satisfied_id_; 77 jmethodID create_example_iterator_id_; 78 }; 79 80 } // namespace jni 81 } // namespace engine 82 } // namespace client 83 } // namespace fcp 84